stmmlib 0.2.0__py3-none-any.whl

stmmlib/__init__.py

@@ -0,0 +1,30 @@
+"""stmmlib — Family-aware spatio-temporal mixed models.
+
+Supports 14 model structures (M1-M14) for Gaussian and Beta families,
+with optional CUDA/GPU acceleration via CuPy.
+"""
+
+from stmmlib.registry import (
+    stmm_available_models,
+    stmm_model_label,
+    stmm_parse_model,
+)
+from stmmlib.stmm import stmm, stmm_normal, stmm_beta
+from stmmlib.output_format import format_output
+from stmmlib._backend import Backend
+
+try:  # single source of truth for the version: the installed package metadata
+    from importlib.metadata import version as _pkg_version
+    __version__ = _pkg_version("stmmlib")
+except Exception:  # source checkout without an install
+    __version__ = "0.2.0"
+__all__ = [
+    "stmm",
+    "stmm_normal",
+    "stmm_beta",
+    "format_output",
+    "stmm_available_models",
+    "stmm_model_label",
+    "stmm_parse_model",
+    "Backend",
+]

stmmlib/_backend.py

@@ -0,0 +1,153 @@
+"""GPU/CPU backend abstraction with auto-detection.
+
+Usage:
+    backend = Backend(gpu="auto")  # auto-detect GPU
+    backend = Backend(gpu=True)    # force GPU (error if unavailable)
+    backend = Backend(gpu=False)   # force CPU
+
+All array operations go through backend.xp (numpy or cupy).
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Literal
+
+
+class Backend:
+    """Array backend that routes computation to CPU or GPU.
+
+    Attributes:
+        xp: Array namespace (numpy or cupy).
+        sparse: Sparse matrix namespace (scipy.sparse or cupyx.scipy.sparse).
+        device: "cpu" or "gpu".
+    """
+
+    def __init__(self, gpu: bool | Literal["auto"] = "auto") -> None:
+        self.xp = None
+        self.sparse = None
+        self.device: Literal["cpu", "gpu"] = "cpu"
+
+        use_gpu = self._resolve_gpu(gpu)
+
+        if use_gpu:
+            try:
+                import cupy as cp  # type: ignore[import-untyped]
+                import cupyx.scipy.sparse as csp  # type: ignore[import-untyped]
+
+                cp.zeros(1, dtype=cp.float64)  # test allocation
+                self.xp = cp  # type: ignore[assignment]
+                self.sparse = csp  # type: ignore[assignment]
+                self.device = "gpu"
+                return
+            except Exception as e:
+                if gpu is True:
+                    raise RuntimeError(f"GPU requested but unavailable: {e}") from e
+                warnings.warn(
+                    f"GPU unavailable ({e}), falling back to CPU.",
+                    stacklevel=2,
+                )
+
+        import numpy as np
+        import scipy.sparse as sp
+
+        self.xp = np  # type: ignore[assignment]
+        self.sparse = sp  # type: ignore[assignment]
+        self.device = "cpu"
+
+    @staticmethod
+    def _resolve_gpu(gpu: bool | Literal["auto"]) -> bool:
+        if isinstance(gpu, str):
+            g = gpu.strip().lower()
+            if g == "auto":
+                try:
+                    import cupy as cp  # type: ignore[import-untyped]
+
+                    cp.zeros(1, dtype=cp.float64)
+                    return True
+                except Exception:
+                    return False
+            if g in ("no", "false", "0", "off", "cpu"):
+                return False
+            if g in ("yes", "true", "1", "on", "gpu"):
+                return True
+        if gpu == "auto":
+            try:
+                import cupy as cp  # type: ignore[import-untyped]
+
+                cp.zeros(1, dtype=cp.float64)
+                return True
+            except Exception:
+                return False
+        return bool(gpu)
+
+    def to_cpu(self, arr):
+        """Move array to CPU. No-op if already on CPU."""
+        if self.device == "gpu":
+            return self.xp.asnumpy(arr)
+        return arr
+
+    def to_gpu(self, arr):
+        """Move array to GPU. No-op if already on CPU."""
+        if self.device == "gpu":
+            return self.xp.asarray(arr)
+        return arr
+
+    def is_sparse(self, m) -> bool:
+        """Check if matrix is a sparse matrix."""
+        if self.device == "gpu":
+            return self.sparse.issparse(m)
+        import scipy.sparse as sp
+
+        return sp.issparse(m)
+
+    def to_sparse(self, m):
+        """Convert dense array to sparse CSR."""
+        if self.is_sparse(m):
+            return m
+        if self.device == "gpu":
+            return self.sparse.csr_matrix(m)
+        import scipy.sparse as sp
+
+        return sp.csr_matrix(m)
+
+    def sparse_kron(self, A, B, format="csr"):
+        """Kronecker product of two (sparse) matrices."""
+        return self.sparse.kron(A, B, format=format)
+
+    def sparse_block_diag(self, mats):
+        """Block-diagonal from list of matrices."""
+        if self.device == "gpu":
+            return self.sparse.block_diag(mats, format="csr")
+        import scipy.sparse as sp
+
+        return sp.block_diag(mats, format="csr")
+
+    def sparse_eye(self, n, dtype=None):
+        """Identity matrix (sparse)."""
+        if dtype is None:
+            dtype = self.xp.float64
+        return self.sparse.eye(n, dtype=dtype, format="csr")
+
+    def solve(self, A, b):
+        """Solve linear system A @ x = b."""
+        return self.xp.linalg.solve(A, b)
+
+    def cholesky(self, A):
+        """Cholesky decomposition. Returns lower-triangular L where A = L @ L.T."""
+        return self.xp.linalg.cholesky(A)
+
+    def slogdet(self, A):
+        """Log-determinant via sign and log-abs-determinant."""
+        return self.xp.linalg.slogdet(A)
+
+    def eigh(self, A):
+        """Eigenvalues and eigenvectors of symmetric matrix."""
+        return self.xp.linalg.eigh(A)
+
+    def pinv(self, A):
+        """Moore-Penrose pseudo-inverse."""
+        return self.xp.linalg.pinv(A)
+
+    def __repr__(self) -> str:
+        return f"Backend(device={self.device!r})"

stmmlib/_utils.py

@@ -0,0 +1,175 @@
+"""Shared matrix utilities used by both Gaussian and Beta backends.
+
+All functions accept a Backend instance and use backend.xp for array ops.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from stmmlib._backend import Backend
+
+
+# ---------------------------------------------------------------------------
+# Positive-definiteness checks
+# ---------------------------------------------------------------------------
+
+
+def is_pd(M, backend: Backend, tol: float = 1e-8) -> bool:
+    """Test if matrix M is symmetric positive-definite via Cholesky attempt."""
+    xp = backend.xp
+    if backend.is_sparse(M):
+        M = M.toarray()
+    M = xp.asarray(M, dtype=xp.float64)
+    # Check approximate symmetry
+    if not xp.allclose(M, M.T, atol=tol):
+        return False
+    try:
+        backend.cholesky(M)
+        return True
+    except Exception:
+        return False
+
+
+def near_pd(M, backend: Backend, eps: float = 1e-6):
+    """Project matrix to nearest positive-definite if needed.
+
+    If M is already PD, returns M unchanged. Otherwise adds small ridge
+    to the diagonal and retries.
+    """
+    xp = backend.xp
+    if backend.is_sparse(M):
+        M = M.toarray()
+    M = xp.asarray(M, dtype=xp.float64)
+
+    if is_pd(M, backend, tol=eps):
+        return M
+
+    # Symmetrize
+    M = 0.5 * (M + M.T)
+
+    # Add increasing ridge until PD
+    n = M.shape[0]
+    for k in range(50):
+        shift = eps * (10.0**k) * xp.eye(n, dtype=xp.float64)
+        M_try = M + shift
+        if is_pd(M_try, backend, tol=eps):
+            return M_try
+
+    warnings.warn(
+        f"Matrix not PD after max ridge iterations. Returning with ridge={eps * 10**49}.",
+        stacklevel=2,
+    )
+    return M + eps * 1e49 * xp.eye(n, dtype=xp.float64)
+
+
+def safe_cholesky(A, backend: Backend):
+    """Cholesky decomposition with PD fallback.
+
+    Returns lower-triangular L such that A = L @ L.T.
+    """
+    xp = backend.xp
+    if backend.is_sparse(A):
+        A = A.toarray()
+    A = xp.asarray(A, dtype=xp.float64)
+    try:
+        return backend.cholesky(A)
+    except Exception:
+        A_pd = near_pd(A, backend)
+        return backend.cholesky(A_pd)
+
+
+def safe_solve(A, b, backend: Backend):
+    """Solve A @ x = b with fallback to pseudo-inverse."""
+    xp = backend.xp
+    if backend.is_sparse(A):
+        A = A.toarray()
+    A = xp.asarray(A, dtype=xp.float64)
+    b = xp.asarray(b, dtype=xp.float64)
+    try:
+        return backend.solve(A, b)
+    except Exception:
+        return backend.pinv(A) @ b
+
+
+def logdet_matrix(M, backend: Backend):
+    """Log-determinant of a matrix. Uses slogdet for numerical stability."""
+    xp = backend.xp
+    if backend.is_sparse(M):
+        M = M.toarray()
+    M = xp.asarray(M, dtype=xp.float64)
+    sign, logabsdet = backend.slogdet(M)
+    if float(sign) < 0:
+        warnings.warn("Negative determinant encountered.", stacklevel=2)
+    return float(logabsdet)
+
+
+# ---------------------------------------------------------------------------
+# Interaction helpers
+# ---------------------------------------------------------------------------
+
+
+def interaction_parse(interaction: str) -> str:
+    """Normalize interaction type string."""
+    mapping = {
+        "none": "none",
+        "iid": "iid",
+        "spatial": "spatial",
+        "temporal": "temporal",
+        "st": "st",
+        "Type I": "iid",
+        "Type II": "spatial",
+        "Type III": "temporal",
+        "Type IV": "st",
+    }
+    return mapping.get(str(interaction).strip().lower(), "none")
+
+
+def interaction_active(interaction: str) -> bool:
+    """True if interaction is not 'none'."""
+    return interaction_parse(interaction) != "none"
+
+
+def row_standardize(W, safe: bool = True):
+    """Row-standardize a spatial weight matrix, island-safe.
+
+    Args:
+        W: (m, m) array-like weight matrix (numpy or scipy sparse).
+        safe: True (default) -> rows with zero sum (islands, no neighbors)
+            are LEFT AS ZERO (no NaN); False -> plain W/row_sum (0/0 -> NaN,
+            the classic island trap).
+
+    Returns:
+        numpy (m, m) row-standardized matrix. Island rows stay zero so that
+        A = I - rho W keeps them iid under the SAR prior (valid: no
+        neighbors -> no spatial shrinkage for that area).
+    """
+    import numpy as np
+    from scipy import sparse as _sp
+    if _sp.issparse(W):
+        W = W.toarray()
+    W = np.asarray(W, dtype=np.float64)
+    rs = W.sum(axis=1, keepdims=True)
+    if safe:
+        return W / np.maximum(rs, 1.0)
+    return W / rs
+
+
+def interaction_param_names(interaction: str) -> list[str]:
+    """Extra variance parameter names for interaction type."""
+    itype = interaction_parse(interaction)
+    if itype == "none":
+        return []
+    if itype == "iid":
+        return ["sigma2_gamma"]
+    if itype == "spatial":
+        return ["sigma2_gamma", "rho_gamma"]
+    if itype == "temporal":
+        return ["sigma2_gamma", "phi_gamma"]
+    if itype == "st":
+        return ["sigma2_gamma", "rho_gamma", "phi_gamma"]
+    return []
+
+

stmmlib/beta/__init__.py

@@ -0,0 +1,17 @@
+"""Beta family backend for spatio-temporal mixed models.
+
+Modules:
+    structure: M1-M14 model registry
+    precision: AR(1) precision, rho/phi bounds, effective phi
+    design: Z matrices, contrast matrices, W_gamma
+    solvers: Newton/BFGS mode solver, Laplace marginal likelihood engine
+    inference: Hessian-based SEs, Wald tests, confidence intervals
+    result: BetastmmResult container
+    fit: stmm_beta(), betaml() top-level fitting
+    diagnostics: Residuals, ACF, Dunn-Smyth, Moran panel, Ljung-Box
+    criteria: conditional AIC / conditional R-squared for flat models
+    report: diagnostics report and convergence certificate
+    lrt_fixed: restricted-refit likelihood ratio test for fixed effects
+    config: numerical bounds and tolerances (single source of truth)
+    moran_panel, plot_data, gpu_solvers, simulation, diagnose: helpers
+"""

stmmlib/beta/config.py

@@ -0,0 +1,108 @@
+"""Central numeric constants for Beta STMM fitting.
+
+Single source of truth for bounds, tolerances, and accuracy presets
+previously scattered as literals across solvers/precision/design/criteria.
+All defaults reproduce the historical behavior exactly.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Auto weak-prior constants (blme-style, Chung et al. 2013): center 6.9 ~
+# phi 1000. These do NOT move the box bounds, only the prior center/spread.
+# With two coupled directions (phi + dep) the weak-prior path can collapse
+# into a degenerate well, so the dep part of the auto prior is off by default.
+PHI_PRIOR_AUTO_MEAN = 6.9
+PHI_PRIOR_AUTO_SD = 2.0
+DEP_PRIOR_AUTO_SD = None
+
+
+@dataclass(frozen=True)
+class BetaFitConfig:
+    """Numeric config for Beta STMM fitting (all defaults = legacy behavior).
+
+    Weak-prior penalty (opt-in MAP, default OFF):
+        phi_prior_mean / phi_prior_sd: Gaussian prior on logphi, penalty
+            0.5*((logphi-mu)/sd)^2. dep_prior_sd: Gaussian prior N(0, sd^2)
+            on raw_dep, penalty 0.5*sum(raw_dep^2)/sd^2. All default None
+            (= unpenalized ML; behavior identical to before).
+        MAP with weak Gaussian prior; precedent Chung et al. 2013 blme
+        (Psychometrika) and brglm2 MPL Jeffreys. Always report penalized
+        alongside the unpenalized solution.
+        Note: passing an explicit config (even a default BetaFitConfig())
+        disables the automatic boundary refit in stmm_beta(); leaving
+        config=None lets stmm_beta auto-refit with a weak prior
+        (PHI_PRIOR_AUTO_MEAN/SD, DEP_PRIOR_AUTO_SD) when the unpenalized
+        fit hits the log_phi/raw_dep boundary.
+        hessian_time_budget 3600.0 = cap 1 jam untuk Hessian (default; cegah
+        hang). None = tanpa batas waktu. Set angka detik lain untuk mengubah cap.
+    """
+
+    phi_bound: float = 0.95
+    # 0.95 is a conventional stationarity margin for the dependence
+    # coefficients: rho = 0.95 corresponds to a half-life of ~13.5 periods.
+    # Do not raise the bound just to escape the boundary.
+    #
+    # Two correlation bounds, not one:
+    #   phi_bound: the AR(1) coefficient bound (phi_T, phi_gamma).
+    #   rho_bound: the spatial correlation bound (SAR rho, gamma rho).
+    # Neither is the Beta precision: that one is logphi_bounds.
+    rho_bound: float = 0.95
+    beta_bounds: tuple = (-20.0, 20.0)
+    logphi_bounds: tuple = (-10.0, 10.0)
+    logs_bounds: tuple = (-10.0, 3.0)
+    dep_margin: float = 0.95
+    grad_tol: float = 1.0
+    grad_slack: float = 100.0
+    # Inner-solve budget for the analytic-gradient path on large panels.
+    # Set from one large panel (the fit jumps inner modes and returns a
+    # garbage gradient with a full budget); it has not been swept across
+    # panel sizes. A smaller cap trades accuracy for stability; the deep
+    # gradient solve below is unaffected. Re-tune before relying on it for
+    # a new model family or a much larger panel.
+    inner_probe_cap: int = 15
+    # Auto weak-prior refit. None keeps the historical rule (the refit runs
+    # only when no explicit config was passed); True always runs it, False
+    # never runs it.
+    auto_prior: bool | None = None
+    # Interim dataset-calibrated prior; always report alongside unpenalized ML (Chung et al. 2013).
+    phi_prior_mean: float | None = None
+    phi_prior_sd: float | None = None
+    dep_prior_sd: float | None = None
+    hessian_time_budget: float | None = 3600.0
+    # Laplace posterior curvature of the random effects at the final mode
+    # (H_b = Z'diag(w)Z + Q), stored as its Cholesky factor for EBP draws.
+    # Opt-in: dense q x q precision Cholesky, one per fit; large for big
+# random-effect designs. Default off.
+    compute_posterior: bool = False
+
+    def __post_init__(self):
+        if (self.phi_prior_mean is None) != (self.phi_prior_sd is None):
+            raise ValueError(
+                "phi_prior_mean and phi_prior_sd must both be set or both None.")
+        if self.phi_prior_sd is not None and self.phi_prior_sd <= 0:
+            raise ValueError("phi_prior_sd must be > 0.")
+        if self.dep_prior_sd is not None and self.dep_prior_sd <= 0:
+            raise ValueError("dep_prior_sd must be > 0.")
+        if self.hessian_time_budget is not None and self.hessian_time_budget <= 0:
+            raise ValueError("hessian_time_budget must be > 0 or None.")
+
+
+PRESET_ACCURACY: dict = {
+    "fast": (1e-5, 1e-5, False, 500),
+    "balanced": (1e-8, 1e-6, True, 500),
+    "medium": (1e-8, 1e-6, False, None),
+    "full": (1e-12, 1e-6, True, None),
+}
+
+
+def resolve_preset(preset: str | None) -> tuple:
+    """Resolve accuracy preset name to (ftol, gtol, hess, cap)."""
+    if preset is None:
+        preset = "full"
+    try:
+        return PRESET_ACCURACY[preset]
+    except KeyError:
+        raise ValueError(
+            f"Unknown preset: {preset!r}. Use 'fast', 'balanced', 'medium' or 'full'.")