nrl-tracker 1.9.1__py3-none-any.whl → 1.10.0__py3-none-any.whl

pytcl/gpu/utils.py

@@ -0,0 +1,582 @@
+"""
+GPU utility functions for array management and device detection.
+
+This module provides utilities for:
+- Checking GPU availability (CUDA via CuPy or Apple Silicon via MLX)
+- Transferring arrays between CPU and GPU
+- Getting the appropriate array module (numpy, cupy, or mlx)
+- Memory management
+- Automatic backend selection based on platform
+
+The module automatically selects the appropriate backend:
+- On Apple Silicon (M1/M2/M3): Uses MLX if available
+- On systems with NVIDIA GPUs: Uses CuPy if available
+- Falls back to CPU (numpy) if no GPU backend is available
+
+Examples
+--------
+>>> from pytcl.gpu.utils import is_gpu_available, to_gpu, to_cpu
+>>> if is_gpu_available():
+...     x_gpu = to_gpu(x_numpy)
+...     # ... perform GPU operations ...
+...     x_cpu = to_cpu(x_gpu)
+"""
+
+import logging
+import platform
+from functools import lru_cache
+from typing import Any, Literal, Union
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+from pytcl.core.optional_deps import is_available
+
+# Module logger
+_logger = logging.getLogger("pytcl.gpu.utils")
+
+# Type alias for arrays that could be numpy, cupy, or mlx
+GPUArray = Any  # Would be cp.ndarray or mx.array if backend is available
+
+# Backend type
+BackendType = Literal["cupy", "mlx", "numpy"]
+
+
+@lru_cache(maxsize=1)
+def is_apple_silicon() -> bool:
+    """
+    Check if running on Apple Silicon (ARM64 Mac).
+
+    Returns
+    -------
+    bool
+        True if running on Apple Silicon (M1, M2, M3, etc.).
+
+    Examples
+    --------
+    >>> from pytcl.gpu.utils import is_apple_silicon
+    >>> if is_apple_silicon():
+    ...     print("Running on Apple Silicon")
+    """
+    return platform.system() == "Darwin" and platform.machine() == "arm64"
+
+
+@lru_cache(maxsize=1)
+def is_mlx_available() -> bool:
+    """
+    Check if MLX acceleration is available (Apple Silicon).
+
+    Returns True if:
+    - Running on Apple Silicon (ARM64 Mac)
+    - MLX is installed
+
+    Returns
+    -------
+    bool
+        True if MLX acceleration is available.
+
+    Examples
+    --------
+    >>> from pytcl.gpu.utils import is_mlx_available
+    >>> if is_mlx_available():
+    ...     print("MLX acceleration enabled")
+    """
+    if not is_apple_silicon():
+        _logger.debug("Not on Apple Silicon, MLX not applicable")
+        return False
+
+    if not is_available("mlx"):
+        _logger.debug("MLX not installed")
+        return False
+
+    try:
+        import mlx.core as mx
+
+        # Verify MLX works by creating a simple array
+        _ = mx.array([1.0, 2.0, 3.0])
+        _logger.info("MLX available on Apple Silicon")
+        return True
+    except Exception as e:
+        _logger.debug("MLX not functional: %s", e)
+        return False
+
+
+@lru_cache(maxsize=1)
+def is_cupy_available() -> bool:
+    """
+    Check if CuPy (CUDA) acceleration is available.
+
+    Returns True if:
+    - CuPy is installed
+    - A CUDA-capable GPU is detected
+    - CUDA runtime is functional
+
+    Returns
+    -------
+    bool
+        True if CuPy acceleration is available.
+    """
+    if not is_available("cupy"):
+        _logger.debug("CuPy not installed")
+        return False
+
+    try:
+        import cupy as cp
+
+        # Try to access a GPU device
+        device = cp.cuda.Device(0)
+        _ = device.compute_capability
+        _logger.info("CuPy available: %s", device.pci_bus_id)
+        return True
+    except Exception as e:
+        _logger.debug("CuPy/CUDA not available: %s", e)
+        return False
+
+
+@lru_cache(maxsize=1)
+def get_backend() -> BackendType:
+    """
+    Get the best available GPU backend for the current platform.
+
+    Priority:
+    1. MLX on Apple Silicon
+    2. CuPy on systems with NVIDIA GPUs
+    3. numpy (CPU fallback)
+
+    Returns
+    -------
+    str
+        One of "mlx", "cupy", or "numpy".
+
+    Examples
+    --------
+    >>> from pytcl.gpu.utils import get_backend
+    >>> backend = get_backend()
+    >>> print(f"Using {backend} backend")
+    """
+    if is_apple_silicon() and is_mlx_available():
+        return "mlx"
+    elif is_cupy_available():
+        return "cupy"
+    else:
+        return "numpy"
+
+
+@lru_cache(maxsize=1)
+def is_gpu_available() -> bool:
+    """
+    Check if GPU acceleration is available.
+
+    Returns True if either:
+    - MLX is available (Apple Silicon)
+    - CuPy is available with a CUDA GPU
+
+    Returns
+    -------
+    bool
+        True if GPU acceleration is available.
+
+    Examples
+    --------
+    >>> from pytcl.gpu.utils import is_gpu_available
+    >>> if is_gpu_available():
+    ...     print("GPU acceleration enabled")
+    ... else:
+    ...     print("Falling back to CPU")
+
+    Notes
+    -----
+    The result is cached after the first call for performance.
+    Use `get_backend()` to determine which backend is being used.
+    """
+    return is_mlx_available() or is_cupy_available()
+
+
+def get_array_module(arr: ArrayLike) -> Any:
+    """
+    Get the array module (numpy, cupy, or mlx.core) for the given array.
+
+    This function enables writing code that works with numpy, cupy, and mlx
+    arrays by returning the appropriate module.
+
+    Parameters
+    ----------
+    arr : array_like
+        Input array (numpy, cupy, or mlx).
+
+    Returns
+    -------
+    module
+        numpy, cupy, or mlx.core module, depending on the input array type.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.utils import get_array_module
+    >>> x = np.array([1, 2, 3])
+    >>> xp = get_array_module(x)
+    >>> xp is np
+    True
+
+    >>> # With CuPy array
+    >>> import cupy as cp
+    >>> x_gpu = cp.array([1, 2, 3])
+    >>> xp = get_array_module(x_gpu)
+    >>> xp is cp
+    True
+
+    >>> # With MLX array
+    >>> import mlx.core as mx
+    >>> x_mlx = mx.array([1, 2, 3])
+    >>> xp = get_array_module(x_mlx)
+    >>> xp.__name__
+    'mlx.core'
+    """
+    # Check for MLX array first
+    if is_available("mlx"):
+        import mlx.core as mx
+
+        if isinstance(arr, mx.array):
+            return mx
+
+    # Check for CuPy array
+    if is_available("cupy"):
+        import cupy as cp
+
+        if isinstance(arr, cp.ndarray):
+            return cp
+
+    return np
+
+
+def to_gpu(arr: ArrayLike, dtype: Any = None, backend: BackendType = None) -> GPUArray:
+    """
+    Transfer an array to GPU memory.
+
+    Automatically selects the best available backend (MLX on Apple Silicon,
+    CuPy on NVIDIA GPUs) unless a specific backend is requested.
+
+    Parameters
+    ----------
+    arr : array_like
+        Input array (typically numpy).
+    dtype : dtype, optional
+        Data type for the GPU array. If None, uses the input dtype.
+    backend : str, optional
+        Specific backend to use ("mlx", "cupy"). If None, auto-selects.
+
+    Returns
+    -------
+    GPUArray
+        Array in GPU memory (cupy.ndarray or mlx.array).
+
+    Raises
+    ------
+    DependencyError
+        If required backend is not installed.
+    RuntimeError
+        If no GPU is available.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.utils import to_gpu, is_gpu_available
+    >>> x = np.array([1.0, 2.0, 3.0])
+    >>> if is_gpu_available():
+    ...     x_gpu = to_gpu(x)
+    ...     print(type(x_gpu).__name__)
+    'ndarray'  # cupy.ndarray or 'array' for mlx
+
+    Notes
+    -----
+    If the input is already a GPU array, it is returned as-is (or converted
+    to the requested dtype).
+    """
+    from pytcl.core.optional_deps import import_optional
+
+    if not is_gpu_available():
+        raise RuntimeError(
+            "No GPU available. Check CUDA installation or MLX availability."
+        )
+
+    # Determine backend
+    if backend is None:
+        backend = get_backend()
+
+    # Use MLX backend
+    if backend == "mlx":
+        mx = import_optional(
+            "mlx.core",
+            package="mlx",
+            extra="gpu-apple",
+            feature="Apple Silicon GPU acceleration",
+        )
+
+        # If already an MLX array
+        if isinstance(arr, mx.array):
+            if dtype is not None:
+                # MLX uses different dtype handling
+                return arr.astype(_numpy_dtype_to_mlx(mx, dtype))
+            return arr
+
+        # Convert to numpy first if needed
+        arr_np = np.asarray(arr)
+        if dtype is not None:
+            arr_np = arr_np.astype(dtype)
+
+        return mx.array(arr_np)
+
+    # Use CuPy backend
+    else:
+        cp = import_optional("cupy", extra="gpu", feature="GPU acceleration")
+
+        # If already a CuPy array
+        if isinstance(arr, cp.ndarray):
+            if dtype is not None and arr.dtype != dtype:
+                return arr.astype(dtype)
+            return arr
+
+        # Convert to numpy first if needed
+        arr_np = np.asarray(arr)
+        if dtype is not None:
+            arr_np = arr_np.astype(dtype)
+
+        return cp.asarray(arr_np)
+
+
+def _numpy_dtype_to_mlx(mx, dtype) -> Any:
+    """Convert numpy dtype to MLX dtype."""
+    dtype_map = {
+        np.float32: mx.float32,
+        np.float64: mx.float32,  # MLX prefers float32
+        np.int32: mx.int32,
+        np.int64: mx.int64,
+        np.bool_: mx.bool_,
+    }
+    if hasattr(dtype, "type"):
+        dtype = dtype.type
+    return dtype_map.get(dtype, mx.float32)
+
+
+def to_cpu(arr: Union[ArrayLike, GPUArray]) -> NDArray[np.floating]:
+    """
+    Transfer an array from GPU to CPU memory.
+
+    Parameters
+    ----------
+    arr : array_like, cupy.ndarray, or mlx.array
+        Input array (numpy, cupy, or mlx).
+
+    Returns
+    -------
+    numpy.ndarray
+        Array in CPU memory.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pytcl.gpu.utils import to_gpu, to_cpu, is_gpu_available
+    >>> x = np.array([1.0, 2.0, 3.0])
+    >>> if is_gpu_available():
+    ...     x_gpu = to_gpu(x)
+    ...     x_cpu = to_cpu(x_gpu)
+    ...     np.allclose(x, x_cpu)
+    True
+
+    Notes
+    -----
+    If the input is already a numpy array, it is returned as-is.
+    """
+    # Already numpy
+    if isinstance(arr, np.ndarray):
+        return arr
+
+    # Check if it's an MLX array
+    if is_available("mlx"):
+        import mlx.core as mx
+
+        if isinstance(arr, mx.array):
+            return np.array(arr)
+
+    # Check if it's a CuPy array
+    if is_available("cupy"):
+        import cupy as cp
+
+        if isinstance(arr, cp.ndarray):
+            return cp.asnumpy(arr)
+
+    # Fallback: convert via numpy
+    return np.asarray(arr)
+
+
+def ensure_gpu_array(
+    arr: ArrayLike,
+    dtype: Any = np.float64,
+    backend: BackendType = None,
+) -> GPUArray:
+    """
+    Ensure an array is on the GPU with the specified dtype.
+
+    Parameters
+    ----------
+    arr : array_like
+        Input array.
+    dtype : dtype
+        Desired data type.
+    backend : str, optional
+        Specific backend to use ("mlx", "cupy"). If None, auto-selects.
+
+    Returns
+    -------
+    GPUArray
+        Array on GPU with specified dtype (cupy.ndarray or mlx.array).
+    """
+    gpu_arr = to_gpu(arr, backend=backend)
+
+    # MLX doesn't support float64 well, use float32
+    if backend == "mlx" or (backend is None and get_backend() == "mlx"):
+        if dtype == np.float64:
+            dtype = np.float32
+
+    if hasattr(gpu_arr, "dtype") and gpu_arr.dtype != dtype:
+        if get_backend() == "mlx":
+            import mlx.core as mx
+
+            gpu_arr = gpu_arr.astype(_numpy_dtype_to_mlx(mx, dtype))
+        else:
+            gpu_arr = gpu_arr.astype(dtype)
+    return gpu_arr
+
+
+def sync_gpu() -> None:
+    """
+    Synchronize GPU operations.
+
+    This blocks until all pending GPU operations are complete.
+    Useful for accurate timing measurements.
+
+    Examples
+    --------
+    >>> import time
+    >>> from pytcl.gpu.utils import sync_gpu, is_gpu_available
+    >>> if is_gpu_available():
+    ...     # ... perform GPU operations ...
+    ...     sync_gpu()  # Wait for completion
+    ...     elapsed = time.time() - start
+    """
+    backend = get_backend()
+
+    if backend == "mlx":
+        import mlx.core as mx
+
+        mx.eval()  # MLX uses lazy evaluation, eval() forces execution
+    elif backend == "cupy":
+        import cupy as cp
+
+        cp.cuda.Stream.null.synchronize()
+
+
+def get_gpu_memory_info() -> dict[str, int]:
+    """
+    Get GPU memory usage information.
+
+    Returns
+    -------
+    dict
+        Dictionary with keys:
+        - 'backend': Backend in use ("mlx", "cupy", or "numpy")
+        - 'free': Free memory in bytes (if available)
+        - 'total': Total memory in bytes (if available)
+        - 'used': Used memory in bytes (if available)
+
+    Examples
+    --------
+    >>> from pytcl.gpu.utils import get_gpu_memory_info, is_gpu_available
+    >>> if is_gpu_available():
+    ...     info = get_gpu_memory_info()
+    ...     print(f"Backend: {info['backend']}")
+    """
+    backend = get_backend()
+
+    if backend == "numpy":
+        return {"backend": "numpy", "free": 0, "total": 0, "used": 0}
+
+    if backend == "mlx":
+        # MLX doesn't expose memory info directly, but we can get device info
+        import mlx.core as mx
+
+        device = mx.default_device()
+        return {
+            "backend": "mlx",
+            "device": str(device),
+            "free": -1,  # Not available
+            "total": -1,  # Not available
+            "used": -1,  # Not available
+        }
+
+    # CuPy backend
+    import cupy as cp
+
+    mempool = cp.get_default_memory_pool()
+    free, total = cp.cuda.Device().mem_info
+
+    return {
+        "backend": "cupy",
+        "free": free,
+        "total": total,
+        "used": total - free,
+        "pool_used": mempool.used_bytes(),
+        "pool_total": mempool.total_bytes(),
+    }
+
+
+def clear_gpu_memory() -> None:
+    """
+    Clear GPU memory pools.
+
+    This frees cached memory blocks held by the GPU backend.
+    Call this when you need to free GPU memory for other operations.
+
+    Examples
+    --------
+    >>> from pytcl.gpu.utils import clear_gpu_memory, is_gpu_available
+    >>> if is_gpu_available():
+    ...     # ... perform GPU operations ...
+    ...     clear_gpu_memory()  # Free cached memory
+    """
+    backend = get_backend()
+
+    if backend == "mlx":
+        import mlx.core as mx
+
+        # MLX has automatic memory management, but we can force a sync
+        mx.eval()
+        # Note: MLX doesn't have explicit memory pool clearing like CuPy
+    elif backend == "cupy":
+        import cupy as cp
+
+        mempool = cp.get_default_memory_pool()
+        mempool.free_all_blocks()
+
+
+__all__ = [
+    # Platform detection
+    "is_apple_silicon",
+    "is_mlx_available",
+    "is_cupy_available",
+    "get_backend",
+    # Availability check
+    "is_gpu_available",
+    # Array operations
+    "get_array_module",
+    "to_gpu",
+    "to_cpu",
+    "ensure_gpu_array",
+    # Synchronization and memory
+    "sync_gpu",
+    "get_gpu_memory_info",
+    "clear_gpu_memory",
+    # Type hints
+    "GPUArray",
+    "BackendType",
+]

pytcl/gravity/clenshaw.py

@@ -106,6 +106,17 @@ def clenshaw_sum_order(
         Sum of C terms weighted by Legendre functions.
     sum_S : float
         Sum of S terms weighted by Legendre functions.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> C = np.zeros((5, 5))
+    >>> S = np.zeros((5, 5))
+    >>> C[2, 0] = 1.0  # Only C20 term
+    >>> cos_theta, sin_theta = np.cos(np.pi/4), np.sin(np.pi/4)
+    >>> sum_C, sum_S = clenshaw_sum_order(0, cos_theta, sin_theta, C, S, 4)
+    >>> isinstance(sum_C, float)
+    True
     """
     # Handle edge case
     if m > n_max:
@@ -187,6 +198,18 @@ def clenshaw_sum_order_derivative(
         Derivative of sum_C with respect to theta.
     dsum_S : float
         Derivative of sum_S with respect to theta.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> C = np.zeros((5, 5))
+    >>> S = np.zeros((5, 5))
+    >>> C[2, 0] = -0.0005  # J2-like term
+    >>> cos_theta, sin_theta = np.cos(np.pi/4), np.sin(np.pi/4)
+    >>> sum_C, sum_S, dsum_C, dsum_S = clenshaw_sum_order_derivative(
+    ...     0, cos_theta, sin_theta, C, S, 4)
+    >>> len([sum_C, sum_S, dsum_C, dsum_S])
+    4
     """
     if m > n_max:
         return 0.0, 0.0, 0.0, 0.0
@@ -298,6 +321,19 @@ def clenshaw_geoid(
             \\sum_{m=0}^{n} P_n^m(\\sin\\phi) (C_{nm}\\cos m\\lambda + S_{nm}\\sin m\\lambda)
 
     The n=0 and n=1 terms are excluded as they represent the reference field.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> C = np.zeros((5, 5))
+    >>> S = np.zeros((5, 5))
+    >>> C[0, 0] = 1.0
+    >>> R = 6.378e6
+    >>> GM = 3.986e14
+    >>> gamma = 9.81
+    >>> N = clenshaw_geoid(0, 0, C, S, R, GM, gamma)
+    >>> isinstance(N, float)
+    True
     """
     if n_max is None:
         n_max = C.shape[0] - 1
@@ -381,6 +417,18 @@ def clenshaw_potential(
     -------
     float
         Gravitational potential in m^2/s^2.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> C = np.zeros((5, 5))
+    >>> S = np.zeros((5, 5))
+    >>> C[0, 0] = 1.0  # Central term only
+    >>> R = 6.378e6
+    >>> GM = 3.986e14
+    >>> V = clenshaw_potential(0, 0, R, C, S, R, GM)
+    >>> abs(V - GM/R) / (GM/R) < 0.01  # ~GM/r for central term
+    True
     """
     if n_max is None:
         n_max = C.shape[0] - 1
@@ -464,6 +512,18 @@ def clenshaw_gravity(
         Northward component of gravity disturbance in m/s^2.
     g_lon : float
         Eastward component of gravity disturbance in m/s^2.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> C = np.zeros((5, 5))
+    >>> S = np.zeros((5, 5))
+    >>> C[0, 0] = 1.0
+    >>> R = 6.378e6
+    >>> GM = 3.986e14
+    >>> g_r, g_lat, g_lon = clenshaw_gravity(0, 0, R, C, S, R, GM)
+    >>> g_r < 0  # Gravity points inward
+    True
     """
     if n_max is None:
         n_max = C.shape[0] - 1