nrl-tracker 1.9.2__py3-none-any.whl → 1.11.0__py3-none-any.whl

pytcl/gpu/matrix_utils.py

@@ -0,0 +1,491 @@
+"""
+GPU matrix utilities for numerical linear algebra.
+
+This module provides GPU-accelerated matrix operations commonly used in
+tracking algorithms, including:
+- Cholesky decomposition
+- QR factorization
+- Matrix inversion and solving
+- Memory pool management
+
+Examples
+--------
+>>> from pytcl.gpu.matrix_utils import gpu_cholesky, gpu_solve
+>>> import numpy as np
+>>>
+>>> # Compute Cholesky decomposition on GPU
+>>> A = np.eye(4) + np.random.randn(4, 4) * 0.1
+>>> A = A @ A.T  # Make positive definite
+>>> L = gpu_cholesky(A)
+>>>
+>>> # Solve linear system
+>>> b = np.random.randn(4)
+>>> x = gpu_solve(A, b)
+"""
+
+import logging
+from contextlib import contextmanager
+from typing import Any, Generator, Optional, Tuple
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+from pytcl.core.optional_deps import import_optional, is_available, requires
+from pytcl.gpu.utils import ensure_gpu_array
+
+# Module logger
+_logger = logging.getLogger("pytcl.gpu.matrix_utils")
+
+
+@requires("cupy", extra="gpu", feature="GPU matrix utilities")
+def gpu_cholesky(A: ArrayLike, lower: bool = True) -> NDArray[np.floating[Any]]:
+    """
+    GPU-accelerated Cholesky decomposition.
+
+    Computes L such that A = L @ L.T (lower=True) or A = U.T @ U (lower=False).
+
+    Parameters
+    ----------
+    A : array_like
+        Symmetric positive definite matrix, shape (n, n) or batch (k, n, n).
+    lower : bool
+        If True, return lower triangular. If False, return upper triangular.
+
+    Returns
+    -------
+    L : ndarray
+        Cholesky factor, same shape as A.
+
+    Raises
+    ------
+    numpy.linalg.LinAlgError
+        If matrix is not positive definite.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.matrix_utils import gpu_cholesky
+    >>> A = np.array([[4, 2], [2, 3]])
+    >>> L = gpu_cholesky(A)
+    >>> np.allclose(L @ L.T, A)
+    True
+    """
+    cp = import_optional("cupy", extra="gpu", feature="GPU matrix utilities")
+
+    A_gpu = ensure_gpu_array(A, dtype=cp.float64)
+
+    L = cp.linalg.cholesky(A_gpu)
+
+    if not lower:
+        if A_gpu.ndim == 2:
+            L = L.T
+        else:
+            L = cp.swapaxes(L, -2, -1)
+
+    return L
+
+
+@requires("cupy", extra="gpu", feature="GPU matrix utilities")
+def gpu_cholesky_safe(
+    A: ArrayLike,
+    lower: bool = True,
+    regularization: float = 1e-10,
+) -> Tuple[NDArray[np.floating[Any]], bool]:
+    """
+    GPU Cholesky decomposition with fallback for non-positive-definite matrices.
+
+    If standard Cholesky fails, adds regularization to diagonal and retries.
+
+    Parameters
+    ----------
+    A : array_like
+        Symmetric matrix, shape (n, n) or batch (k, n, n).
+    lower : bool
+        Return lower (True) or upper (False) triangular factor.
+    regularization : float
+        Amount to add to diagonal if matrix is not positive definite.
+
+    Returns
+    -------
+    L : ndarray
+        Cholesky factor.
+    success : bool
+        True if succeeded without regularization.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.matrix_utils import gpu_cholesky_safe
+    >>> A = np.array([[1, 2], [2, 1]])  # Not positive definite
+    >>> L, success = gpu_cholesky_safe(A)
+    >>> success
+    False
+    """
+    cp = import_optional("cupy", extra="gpu", feature="GPU matrix utilities")
+
+    A_gpu = ensure_gpu_array(A, dtype=cp.float64)
+
+    try:
+        L = cp.linalg.cholesky(A_gpu)
+        success = True
+    except cp.linalg.LinAlgError:
+        # Add regularization
+        if A_gpu.ndim == 2:
+            A_reg = A_gpu + regularization * cp.eye(A_gpu.shape[0], dtype=cp.float64)
+        else:
+            # Batch case
+            n = A_gpu.shape[-1]
+            eye = cp.eye(n, dtype=cp.float64)
+            A_reg = A_gpu + regularization * eye
+
+        L = cp.linalg.cholesky(A_reg)
+        success = False
+        _logger.warning("Cholesky decomposition required regularization")
+
+    if not lower:
+        if A_gpu.ndim == 2:
+            L = L.T
+        else:
+            L = cp.swapaxes(L, -2, -1)
+
+    return L, success
+
+
+@requires("cupy", extra="gpu", feature="GPU matrix utilities")
+def gpu_qr(
+    A: ArrayLike, mode: str = "reduced"
+) -> Tuple[NDArray[np.floating[Any]], NDArray[np.floating[Any]]]:
+    """
+    GPU-accelerated QR decomposition.
+
+    Computes A = Q @ R where Q is orthogonal and R is upper triangular.
+
+    Parameters
+    ----------
+    A : array_like
+        Matrix to decompose, shape (m, n) or batch (k, m, n).
+    mode : str
+        'reduced' (default) or 'complete'.
+
+    Returns
+    -------
+    Q : ndarray
+        Orthogonal matrix.
+    R : ndarray
+        Upper triangular matrix.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.matrix_utils import gpu_qr
+    >>> A = np.random.randn(4, 3)
+    >>> Q, R = gpu_qr(A)
+    >>> np.allclose(Q @ R, A)
+    True
+    """
+    cp = import_optional("cupy", extra="gpu", feature="GPU matrix utilities")
+
+    A_gpu = ensure_gpu_array(A, dtype=cp.float64)
+    Q, R = cp.linalg.qr(A_gpu, mode=mode)
+
+    return Q, R
+
+
+@requires("cupy", extra="gpu", feature="GPU matrix utilities")
+def gpu_solve(A: ArrayLike, b: ArrayLike) -> NDArray[np.floating[Any]]:
+    """
+    GPU-accelerated linear system solve.
+
+    Solves A @ x = b for x.
+
+    Parameters
+    ----------
+    A : array_like
+        Coefficient matrix, shape (n, n) or batch (k, n, n).
+    b : array_like
+        Right-hand side, shape (n,) or (n, m) or batch (k, n).
+
+    Returns
+    -------
+    x : ndarray
+        Solution vector/matrix.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.matrix_utils import gpu_solve
+    >>> A = np.array([[3, 1], [1, 2]])
+    >>> b = np.array([9, 8])
+    >>> x = gpu_solve(A, b)
+    >>> np.allclose(A @ x, b)
+    True
+    """
+    cp = import_optional("cupy", extra="gpu", feature="GPU matrix utilities")
+
+    A_gpu = ensure_gpu_array(A, dtype=cp.float64)
+    b_gpu = ensure_gpu_array(b, dtype=cp.float64)
+
+    x = cp.linalg.solve(A_gpu, b_gpu)
+
+    return x
+
+
+@requires("cupy", extra="gpu", feature="GPU matrix utilities")
+def gpu_inv(A: ArrayLike) -> NDArray[np.floating[Any]]:
+    """
+    GPU-accelerated matrix inversion.
+
+    Parameters
+    ----------
+    A : array_like
+        Matrix to invert, shape (n, n) or batch (k, n, n).
+
+    Returns
+    -------
+    A_inv : ndarray
+        Inverse matrix.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.matrix_utils import gpu_inv
+    >>> A = np.array([[1, 2], [3, 4]])
+    >>> A_inv = gpu_inv(A)
+    >>> np.allclose(A @ A_inv, np.eye(2))
+    True
+    """
+    cp = import_optional("cupy", extra="gpu", feature="GPU matrix utilities")
+
+    A_gpu = ensure_gpu_array(A, dtype=cp.float64)
+    A_inv = cp.linalg.inv(A_gpu)
+
+    return A_inv
+
+
+@requires("cupy", extra="gpu", feature="GPU matrix utilities")
+def gpu_eigh(
+    A: ArrayLike,
+) -> Tuple[NDArray[np.floating[Any]], NDArray[np.floating[Any]]]:
+    """
+    GPU-accelerated eigendecomposition for symmetric matrices.
+
+    Computes eigenvalues and eigenvectors of symmetric matrix A.
+
+    Parameters
+    ----------
+    A : array_like
+        Symmetric matrix, shape (n, n) or batch (k, n, n).
+
+    Returns
+    -------
+    eigenvalues : ndarray
+        Eigenvalues in ascending order.
+    eigenvectors : ndarray
+        Corresponding eigenvectors as columns.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.matrix_utils import gpu_eigh
+    >>> A = np.array([[2, 1], [1, 2]])
+    >>> eigvals, eigvecs = gpu_eigh(A)
+    >>> eigvals
+    array([1., 3.])
+    """
+    cp = import_optional("cupy", extra="gpu", feature="GPU matrix utilities")
+
+    A_gpu = ensure_gpu_array(A, dtype=cp.float64)
+    eigvals, eigvecs = cp.linalg.eigh(A_gpu)
+
+    return eigvals, eigvecs
+
+
+@requires("cupy", extra="gpu", feature="GPU matrix utilities")
+def gpu_matrix_sqrt(A: ArrayLike) -> NDArray[np.floating[Any]]:
+    """
+    GPU-accelerated matrix square root for positive definite matrices.
+
+    Computes S such that S @ S = A using eigendecomposition.
+
+    Parameters
+    ----------
+    A : array_like
+        Symmetric positive definite matrix.
+
+    Returns
+    -------
+    S : ndarray
+        Matrix square root.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.matrix_utils import gpu_matrix_sqrt
+    >>> A = np.array([[4, 0], [0, 9]])
+    >>> S = gpu_matrix_sqrt(A)
+    >>> np.allclose(S @ S, A)
+    True
+    """
+    cp = import_optional("cupy", extra="gpu", feature="GPU matrix utilities")
+
+    A_gpu = ensure_gpu_array(A, dtype=cp.float64)
+
+    # Eigendecomposition
+    eigvals, eigvecs = cp.linalg.eigh(A_gpu)
+
+    # Ensure non-negative eigenvalues
+    eigvals = cp.maximum(eigvals, 0)
+
+    # Compute sqrt
+    sqrt_eigvals = cp.sqrt(eigvals)
+
+    # Reconstruct: S = V @ diag(sqrt(lambda)) @ V'
+    if A_gpu.ndim == 2:
+        S = eigvecs @ cp.diag(sqrt_eigvals) @ eigvecs.T
+    else:
+        # Batch case
+        S = cp.einsum("...ij,...j,...kj->...ik", eigvecs, sqrt_eigvals, eigvecs)
+
+    return S
+
+
+class MemoryPool:
+    """
+    GPU memory pool manager for efficient memory allocation.
+
+    This class provides convenient access to CuPy's memory pool
+    with additional monitoring and management utilities.
+
+    Examples
+    --------
+    >>> from pytcl.gpu.matrix_utils import MemoryPool
+    >>> pool = MemoryPool()
+    >>> print(pool.get_stats())
+    {'used': 0, 'total': 0, 'free': ...}
+    >>>
+    >>> # Allocate some arrays
+    >>> import cupy as cp
+    >>> x = cp.zeros((1000, 1000))
+    >>> print(pool.get_stats())
+    {'used': 8000000, ...}
+    >>>
+    >>> # Free cached memory
+    >>> pool.free_all()
+    """
+
+    def __init__(self) -> None:
+        """Initialize memory pool manager."""
+        if not is_available("cupy"):
+            _logger.warning("CuPy not available, MemoryPool is a no-op")
+            self._pool = None
+            self._pinned_pool = None
+        else:
+            import cupy as cp
+
+            self._pool = cp.get_default_memory_pool()
+            self._pinned_pool = cp.get_default_pinned_memory_pool()
+
+    def get_stats(self) -> dict[str, int]:
+        """
+        Get memory pool statistics.
+
+        Returns
+        -------
+        stats : dict
+            Dictionary with 'used', 'total', and 'free' bytes.
+        """
+        if self._pool is None:
+            return {"used": 0, "total": 0, "free": 0}
+
+        import cupy as cp
+
+        free, total = cp.cuda.Device().mem_info
+
+        return {
+            "used": self._pool.used_bytes(),
+            "total": self._pool.total_bytes(),
+            "free": free,
+            "device_total": total,
+        }
+
+    def free_all(self) -> None:
+        """Free all cached memory blocks."""
+        if self._pool is not None:
+            self._pool.free_all_blocks()
+        if self._pinned_pool is not None:
+            self._pinned_pool.free_all_blocks()
+
+    def set_limit(self, limit: Optional[int] = None) -> None:
+        """
+        Set memory pool limit.
+
+        Parameters
+        ----------
+        limit : int or None
+            Maximum bytes to allocate. None for no limit.
+        """
+        if self._pool is not None:
+            if limit is None:
+                self._pool.set_limit(size=0)  # 0 means no limit
+            else:
+                self._pool.set_limit(size=limit)
+
+    @contextmanager
+    def limit_memory(self, max_bytes: int) -> Generator[None, None, None]:
+        """
+        Context manager for temporary memory limit.
+
+        Parameters
+        ----------
+        max_bytes : int
+            Maximum bytes allowed during context.
+
+        Examples
+        --------
+        >>> pool = MemoryPool()
+        >>> with pool.limit_memory(1e9):  # 1GB limit
+        ...     # Operations here have limited memory
+        ...     pass
+        """
+        if self._pool is None:
+            yield
+            return
+
+        old_limit = self._pool.get_limit()
+        self._pool.set_limit(size=max_bytes)
+        try:
+            yield
+        finally:
+            self._pool.set_limit(size=old_limit)
+
+
+# Global memory pool instance
+_memory_pool: Optional[MemoryPool] = None
+
+
+def get_memory_pool() -> MemoryPool:
+    """
+    Get the global GPU memory pool manager.
+
+    Returns
+    -------
+    pool : MemoryPool
+        Global memory pool instance.
+    """
+    global _memory_pool
+    if _memory_pool is None:
+        _memory_pool = MemoryPool()
+    return _memory_pool
+
+
+__all__ = [
+    "gpu_cholesky",
+    "gpu_cholesky_safe",
+    "gpu_qr",
+    "gpu_solve",
+    "gpu_inv",
+    "gpu_eigh",
+    "gpu_matrix_sqrt",
+    "MemoryPool",
+    "get_memory_pool",
+]