python-gls 0.1.0__py3-none-any.whl

python_gls/correlation/spatial.py

@@ -0,0 +1,190 @@
+"""Spatial correlation structures.
+
+Implements isotropic spatial correlation functions parameterized by
+range and optional nugget effect.
+"""
+
+import numpy as np
+from numpy.typing import NDArray
+
+from python_gls.correlation.base import CorStruct
+
+
+class _SpatialCorStruct(CorStruct):
+    """Base class for spatial correlation structures.
+
+    All spatial structures share:
+    - A range parameter controlling decay
+    - An optional nugget parameter (discontinuity at distance 0)
+    - Distance matrices stored per group
+
+    Parameters
+    ----------
+    range_param : float, optional
+        Initial range parameter (> 0).
+    nugget : bool
+        Whether to include a nugget effect.
+    """
+
+    def __init__(self, range_param: float | None = None, nugget: bool = False):
+        super().__init__()
+        if range_param is not None:
+            if not isinstance(range_param, (int, float)):
+                raise TypeError(
+                    f"range_param must be a number, got {type(range_param).__name__}"
+                )
+            if range_param <= 0:
+                raise ValueError(
+                    f"range_param must be positive, got {range_param}"
+                )
+        if not isinstance(nugget, bool):
+            raise TypeError(f"nugget must be a boolean, got {type(nugget).__name__}")
+        self._nugget = nugget
+        self._distances: dict[int, NDArray] = {}
+        if range_param is not None:
+            if nugget:
+                self._params = np.array([float(range_param), 0.0])
+            else:
+                self._params = np.array([float(range_param)])
+
+    @property
+    def n_params(self) -> int:
+        return 2 if self._nugget else 1
+
+    def set_distances(self, group_id: int, dist_matrix: NDArray) -> None:
+        """Set the distance matrix for a group."""
+        self._distances[group_id] = np.asarray(dist_matrix, dtype=float)
+
+    def set_coordinates(self, group_id: int, coords: NDArray) -> None:
+        """Set coordinates for a group; distances computed automatically."""
+        coords = np.asarray(coords, dtype=float)
+        if coords.ndim == 1:
+            coords = coords[:, None]
+        from scipy.spatial.distance import cdist
+        self._distances[group_id] = cdist(coords, coords)
+
+    def _correlation_function(self, d: NDArray, range_param: float) -> NDArray:
+        """Compute correlation from distances. Override in subclasses."""
+        raise NotImplementedError
+
+    def get_correlation_matrix(self, group_size: int, **kwargs) -> NDArray:
+        if self._params is None:
+            return np.eye(group_size)
+
+        range_param = self._params[0]
+        group_id = kwargs.get("group_id", None)
+
+        if group_id is not None and group_id in self._distances:
+            dist = self._distances[group_id]
+        else:
+            # Default: unit-spaced
+            idx = np.arange(group_size, dtype=float)
+            dist = np.abs(idx[:, None] - idx[None, :])
+
+        R = self._correlation_function(dist, range_param)
+
+        if self._nugget and len(self._params) > 1:
+            nug = 1 / (1 + np.exp(-self._params[1]))  # sigmoid to (0, 1)
+            R = (1 - nug) * R
+            np.fill_diagonal(R, 1.0)
+
+        return R
+
+    def _get_init_params(self, residuals_by_group: list[NDArray]) -> NDArray:
+        # Heuristic: set range to median distance
+        if self._distances:
+            all_dists = []
+            for d in self._distances.values():
+                mask = np.triu(np.ones(d.shape, dtype=bool), k=1)
+                all_dists.extend(d[mask].tolist())
+            if all_dists:
+                range_init = np.median(all_dists)
+            else:
+                range_init = 1.0
+        else:
+            range_init = 1.0
+        if self._nugget:
+            return np.array([range_init, 0.0])
+        return np.array([range_init])
+
+    def _params_to_unconstrained(self, params: NDArray) -> NDArray:
+        # range > 0: use log transform
+        u = np.zeros_like(params)
+        u[0] = np.log(max(params[0], 1e-10))
+        if self._nugget and len(params) > 1:
+            u[1] = params[1]  # already unconstrained (sigmoid applied in get_corr)
+        return u
+
+    def _unconstrained_to_params(self, uparams: NDArray) -> NDArray:
+        p = np.zeros_like(uparams)
+        p[0] = np.exp(uparams[0])
+        if self._nugget and len(uparams) > 1:
+            p[1] = uparams[1]
+        return p
+
+
+class CorExp(_SpatialCorStruct):
+    """Exponential spatial correlation.
+
+    R(d) = exp(-d / range).
+
+    Equivalent to R's `corExp()`.
+    """
+
+    def _correlation_function(self, d: NDArray, range_param: float) -> NDArray:
+        return np.exp(-d / range_param)
+
+
+class CorGaus(_SpatialCorStruct):
+    """Gaussian spatial correlation.
+
+    R(d) = exp(-(d/range)^2).
+
+    Equivalent to R's `corGaus()`.
+    """
+
+    def _correlation_function(self, d: NDArray, range_param: float) -> NDArray:
+        return np.exp(-(d / range_param) ** 2)
+
+
+class CorLin(_SpatialCorStruct):
+    """Linear spatial correlation.
+
+    R(d) = max(1 - d/range, 0).
+
+    Equivalent to R's `corLin()`.
+    """
+
+    def _correlation_function(self, d: NDArray, range_param: float) -> NDArray:
+        return np.maximum(1 - d / range_param, 0)
+
+
+class CorRatio(_SpatialCorStruct):
+    """Rational quadratic spatial correlation.
+
+    R(d) = 1 / (1 + (d/range)^2).
+
+    Equivalent to R's `corRatio()`.
+    """
+
+    def _correlation_function(self, d: NDArray, range_param: float) -> NDArray:
+        return 1.0 / (1.0 + (d / range_param) ** 2)
+
+
+class CorSpher(_SpatialCorStruct):
+    """Spherical spatial correlation.
+
+    R(d) = 1 - 1.5*(d/range) + 0.5*(d/range)^3  for d < range
+    R(d) = 0  for d >= range
+
+    Equivalent to R's `corSpher()`.
+    """
+
+    def _correlation_function(self, d: NDArray, range_param: float) -> NDArray:
+        ratio = d / range_param
+        R = np.where(
+            ratio < 1,
+            1 - 1.5 * ratio + 0.5 * ratio ** 3,
+            0.0,
+        )
+        return R

python_gls/correlation/symm.py

@@ -0,0 +1,85 @@
+"""Unstructured (general symmetric) correlation structure."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+from python_gls.correlation.base import CorStruct
+from python_gls._parametrization import (
+    unconstrained_to_corr,
+    corr_to_unconstrained,
+)
+
+
+class CorSymm(CorStruct):
+    """Unstructured (general symmetric) correlation.
+
+    Estimates all d(d-1)/2 unique off-diagonal correlations freely.
+    Uses spherical parametrization to ensure positive-definiteness.
+
+    Equivalent to R's `corSymm()`.
+
+    Parameters
+    ----------
+    dim : int or None
+        Dimension of the correlation matrix. If None, inferred from data.
+    """
+
+    def __init__(self, dim: int | None = None):
+        super().__init__()
+        if dim is not None:
+            if not isinstance(dim, int):
+                raise TypeError(f"dim must be an integer, got {type(dim).__name__}")
+            if dim < 2:
+                raise ValueError(f"dim must be >= 2 for a correlation matrix, got {dim}")
+        self._dim = dim
+
+    @property
+    def n_params(self) -> int:
+        if self._dim is None:
+            raise ValueError("Dimension not set. Call initialize() first or pass dim=.")
+        return self._dim * (self._dim - 1) // 2
+
+    def get_correlation_matrix(self, group_size: int, **kwargs) -> NDArray:
+        """Build unstructured correlation matrix from current parameters."""
+        if self._params is None:
+            return np.eye(group_size)
+        return unconstrained_to_corr(self._params, group_size)
+
+    def _get_init_params(self, residuals_by_group: list[NDArray]) -> NDArray:
+        """Initialize from sample correlation of residuals."""
+        # Determine dimension from data
+        sizes = [len(r) for r in residuals_by_group]
+        if len(set(sizes)) != 1:
+            d = max(sizes)
+        else:
+            d = sizes[0]
+        self._dim = d
+
+        # Compute sample correlation from residuals
+        # Stack residuals into matrix (n_groups x d)
+        equal_groups = [r for r in residuals_by_group if len(r) == d]
+        if len(equal_groups) >= 2:
+            resid_mat = np.vstack(equal_groups)
+            if resid_mat.shape[0] > d:
+                R_sample = np.corrcoef(resid_mat.T)
+            else:
+                R_sample = np.eye(d)
+        else:
+            R_sample = np.eye(d)
+
+        # Ensure positive-definiteness
+        eigvals = np.linalg.eigvalsh(R_sample)
+        if np.min(eigvals) < 1e-6:
+            R_sample = R_sample + (1e-6 - np.min(eigvals) + 1e-6) * np.eye(d)
+            # Renormalize to correlation
+            d_inv = np.diag(1.0 / np.sqrt(np.diag(R_sample)))
+            R_sample = d_inv @ R_sample @ d_inv
+
+        return corr_to_unconstrained(R_sample)
+
+    def _params_to_unconstrained(self, params: NDArray) -> NDArray:
+        # Params are already unconstrained (spherical parametrization)
+        return params.copy()
+
+    def _unconstrained_to_params(self, uparams: NDArray) -> NDArray:
+        return uparams.copy()

python_gls/likelihood.py

@@ -0,0 +1,302 @@
+"""ML and REML log-likelihood functions for GLS estimation.
+
+Implements profile log-likelihood where fixed effects (beta) are profiled out,
+leaving only correlation and variance parameters to be optimized.
+"""
+
+import warnings
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def _build_omega_block(
+    corr_matrix: NDArray,
+    var_weights: NDArray,
+    sigma2: float,
+) -> NDArray:
+    """Build covariance block: sigma^2 * A^{1/2} R A^{1/2}.
+
+    Parameters
+    ----------
+    corr_matrix : (m, m) correlation matrix for this group.
+    var_weights : (m,) variance weights for this group (standard deviations).
+    sigma2 : scalar residual variance.
+
+    Returns
+    -------
+    Omega block of shape (m, m).
+    """
+    A_half = np.diag(var_weights)
+    return sigma2 * A_half @ corr_matrix @ A_half
+
+
+def _build_omega_inv_block(
+    corr_matrix: NDArray,
+    var_weights: NDArray,
+) -> NDArray:
+    """Build Omega^{-1} block (up to 1/sigma^2 scaling).
+
+    Returns (1/sigma^2) * A^{-1/2} R^{-1} A^{-1/2}.
+    We drop the sigma^2 factor since it cancels in the profile likelihood.
+    """
+    # Guard against zero/near-zero weights
+    safe_weights = np.where(np.abs(var_weights) < 1e-15, 1e-15, var_weights)
+    A_inv_half = np.diag(1.0 / safe_weights)
+    R_inv = np.linalg.solve(corr_matrix, np.eye(corr_matrix.shape[0]))
+    return A_inv_half @ R_inv @ A_inv_half
+
+
+def _safe_log_weights(var_weights: NDArray) -> float:
+    """Compute sum of log(weights) with protection against zero/negative."""
+    safe_weights = np.maximum(np.abs(var_weights), 1e-300)
+    return float(np.sum(np.log(safe_weights)))
+
+
+def profile_loglik_ml(
+    X_groups: list[NDArray],
+    y_groups: list[NDArray],
+    corr_matrices: list[NDArray],
+    var_weights_groups: list[NDArray],
+    nobs: int,
+) -> float:
+    """Profile log-likelihood under ML estimation.
+
+    Beta and sigma^2 are profiled out. The returned value is the
+    concentrated log-likelihood as a function of correlation and variance
+    parameters only.
+
+    Parameters
+    ----------
+    X_groups : list of (m_g, k) design matrices per group.
+    y_groups : list of (m_g,) response vectors per group.
+    corr_matrices : list of (m_g, m_g) correlation matrices per group.
+    var_weights_groups : list of (m_g,) variance weight vectors per group.
+    nobs : int, total number of observations.
+
+    Returns
+    -------
+    float : profile log-likelihood value.
+    """
+    k = X_groups[0].shape[1]
+    N = nobs
+
+    # Accumulate X'Omega^{-1}X and X'Omega^{-1}y and log|Omega| across groups
+    XtOiX = np.zeros((k, k))
+    XtOiy = np.zeros(k)
+    log_det_omega = 0.0
+
+    for Xg, yg, Rg, wg in zip(X_groups, y_groups, corr_matrices, var_weights_groups):
+        Omega_inv = _build_omega_inv_block(Rg, wg)
+        XtOiX += Xg.T @ Omega_inv @ Xg
+        XtOiy += Xg.T @ Omega_inv @ yg
+
+        # log|Omega_g| = log|R_g| + 2*sum(log(w_g)) (sigma^2 factor added later)
+        sign, logdet_R = np.linalg.slogdet(Rg)
+        if sign <= 0:
+            return -np.inf
+        log_det_omega += logdet_R + 2 * _safe_log_weights(wg)
+
+    # Profile beta: beta_hat = (X'Omega^{-1}X)^{-1} X'Omega^{-1}y
+    try:
+        beta_hat = np.linalg.solve(XtOiX, XtOiy)
+    except np.linalg.LinAlgError:
+        return -np.inf
+
+    # Profile sigma^2: sigma^2_hat = (1/N) * sum_g (y_g - X_g beta)' Omega_inv_g (y_g - X_g beta)
+    rss_weighted = 0.0
+    for Xg, yg, Rg, wg in zip(X_groups, y_groups, corr_matrices, var_weights_groups):
+        Omega_inv = _build_omega_inv_block(Rg, wg)
+        resid = yg - Xg @ beta_hat
+        rss_weighted += resid @ Omega_inv @ resid
+
+    sigma2_hat = rss_weighted / N
+
+    if sigma2_hat <= 0:
+        return -np.inf
+
+    # Log-likelihood: -N/2 * log(2*pi) - N/2 * log(sigma^2) - 1/2 * log|Omega/sigma^2| - N/2
+    # where log|Omega| = N*log(sigma^2) + log_det_omega (without sigma^2)
+    loglik = (
+        -0.5 * N * np.log(2 * np.pi)
+        - 0.5 * N * np.log(sigma2_hat)
+        - 0.5 * log_det_omega
+        - 0.5 * N
+    )
+
+    if np.isnan(loglik) or np.isinf(loglik):
+        return -np.inf
+
+    return loglik
+
+
+def profile_loglik_reml(
+    X_groups: list[NDArray],
+    y_groups: list[NDArray],
+    corr_matrices: list[NDArray],
+    var_weights_groups: list[NDArray],
+    nobs: int,
+) -> float:
+    """Profile log-likelihood under REML estimation.
+
+    Like ML, but integrates out the fixed effects for unbiased variance
+    estimation. The REML adjustment adds -0.5 * log|X'Omega^{-1}X|.
+
+    Parameters
+    ----------
+    X_groups : list of (m_g, k) design matrices per group.
+    y_groups : list of (m_g,) response vectors per group.
+    corr_matrices : list of (m_g, m_g) correlation matrices per group.
+    var_weights_groups : list of (m_g,) variance weight vectors per group.
+    nobs : int, total number of observations.
+
+    Returns
+    -------
+    float : profile REML log-likelihood value.
+    """
+    k = X_groups[0].shape[1]
+    N = nobs
+
+    XtOiX = np.zeros((k, k))
+    XtOiy = np.zeros(k)
+    log_det_omega = 0.0
+
+    for Xg, yg, Rg, wg in zip(X_groups, y_groups, corr_matrices, var_weights_groups):
+        Omega_inv = _build_omega_inv_block(Rg, wg)
+        XtOiX += Xg.T @ Omega_inv @ Xg
+        XtOiy += Xg.T @ Omega_inv @ yg
+
+        sign, logdet_R = np.linalg.slogdet(Rg)
+        if sign <= 0:
+            return -np.inf
+        log_det_omega += logdet_R + 2 * _safe_log_weights(wg)
+
+    try:
+        beta_hat = np.linalg.solve(XtOiX, XtOiy)
+    except np.linalg.LinAlgError:
+        return -np.inf
+
+    # REML uses N-k for sigma^2
+    rss_weighted = 0.0
+    for Xg, yg, Rg, wg in zip(X_groups, y_groups, corr_matrices, var_weights_groups):
+        Omega_inv = _build_omega_inv_block(Rg, wg)
+        resid = yg - Xg @ beta_hat
+        rss_weighted += resid @ Omega_inv @ resid
+
+    N_reml = N - k
+    if N_reml <= 0:
+        return -np.inf
+    sigma2_hat = rss_weighted / N_reml
+
+    if sigma2_hat <= 0:
+        return -np.inf
+
+    # REML log-likelihood
+    sign_xtox, logdet_xtox = np.linalg.slogdet(XtOiX)
+    if sign_xtox <= 0:
+        return -np.inf
+
+    loglik = (
+        -0.5 * N_reml * np.log(2 * np.pi)
+        - 0.5 * N_reml * np.log(sigma2_hat)
+        - 0.5 * log_det_omega
+        - 0.5 * logdet_xtox
+        - 0.5 * N_reml
+    )
+
+    if np.isnan(loglik) or np.isinf(loglik):
+        return -np.inf
+
+    return loglik
+
+
+def compute_gls_estimates(
+    X_groups: list[NDArray],
+    y_groups: list[NDArray],
+    corr_matrices: list[NDArray],
+    var_weights_groups: list[NDArray],
+    nobs: int,
+    method: str = "REML",
+) -> tuple[NDArray, NDArray, float, float]:
+    """Compute GLS beta, covariance, sigma^2, and log-likelihood.
+
+    Given the estimated correlation and variance parameters, compute the
+    final GLS estimates.
+
+    Returns
+    -------
+    beta_hat : (k,) estimated coefficients
+    cov_beta : (k, k) covariance of beta estimates
+    sigma2_hat : estimated residual variance
+    loglik : log-likelihood at these estimates
+    """
+    k = X_groups[0].shape[1]
+    N = nobs
+
+    XtOiX = np.zeros((k, k))
+    XtOiy = np.zeros(k)
+    log_det_omega = 0.0
+
+    for Xg, yg, Rg, wg in zip(X_groups, y_groups, corr_matrices, var_weights_groups):
+        Omega_inv = _build_omega_inv_block(Rg, wg)
+        XtOiX += Xg.T @ Omega_inv @ Xg
+        XtOiy += Xg.T @ Omega_inv @ yg
+
+        sign, logdet_R = np.linalg.slogdet(Rg)
+        log_det_omega += logdet_R + 2 * _safe_log_weights(wg)
+
+    try:
+        beta_hat = np.linalg.solve(XtOiX, XtOiy)
+    except np.linalg.LinAlgError as e:
+        raise np.linalg.LinAlgError(
+            "X'Omega^{-1}X is singular. This may indicate perfect collinearity "
+            "in the design matrix or a degenerate correlation structure."
+        ) from e
+
+    rss_weighted = 0.0
+    for Xg, yg, Rg, wg in zip(X_groups, y_groups, corr_matrices, var_weights_groups):
+        Omega_inv = _build_omega_inv_block(Rg, wg)
+        resid = yg - Xg @ beta_hat
+        rss_weighted += resid @ Omega_inv @ resid
+
+    if method == "REML":
+        denom = N - k
+        if denom <= 0:
+            warnings.warn(
+                f"N - k = {denom} <= 0 for REML. Using ML denominator (N={N}) instead.",
+                stacklevel=2,
+            )
+            denom = N
+        sigma2_hat = rss_weighted / denom
+    else:
+        sigma2_hat = rss_weighted / N
+
+    # Covariance of beta: sigma^2 * (X'Omega^{-1}X)^{-1}
+    try:
+        cov_beta = sigma2_hat * np.linalg.inv(XtOiX)
+    except np.linalg.LinAlgError as e:
+        raise np.linalg.LinAlgError(
+            "Failed to invert X'Omega^{-1}X for covariance estimation. "
+            "The design matrix may be singular or near-singular."
+        ) from e
+
+    # Warn if covariance has issues
+    cov_diag = np.diag(cov_beta)
+    if np.any(cov_diag < 0):
+        warnings.warn(
+            "Some variance estimates are negative, indicating numerical instability. "
+            "Standard errors may be unreliable.",
+            stacklevel=2,
+        )
+
+    # Compute log-likelihood at these estimates
+    if method == "REML":
+        loglik = profile_loglik_reml(
+            X_groups, y_groups, corr_matrices, var_weights_groups, nobs
+        )
+    else:
+        loglik = profile_loglik_ml(
+            X_groups, y_groups, corr_matrices, var_weights_groups, nobs
+        )
+
+    return beta_hat, cov_beta, sigma2_hat, loglik