nrl-tracker 1.1.3__py3-none-any.whl → 1.3.0__py3-none-any.whl

pytcl/gravity/egm.py

@@ -21,6 +21,7 @@ References
        https://earth-info.nga.mil/
 """
 
+import logging
 import os
 from functools import lru_cache
 from pathlib import Path
@@ -32,6 +33,9 @@ from numpy.typing import NDArray
 from .clenshaw import clenshaw_gravity, clenshaw_potential
 from .models import WGS84, normal_gravity_somigliana
 
+# Module logger
+_logger = logging.getLogger("pytcl.gravity.egm")
+
 
 class EGMCoefficients(NamedTuple):
     """Earth Gravitational Model coefficients.
@@ -317,6 +321,8 @@ def _load_coefficients_cached(
     data_dir = get_data_dir()
     filepath = data_dir / f"{model}.cof"
 
+    _logger.debug("Loading %s coefficients from %s", model, filepath)
+
     if not filepath.exists():
         raise FileNotFoundError(
             f"Coefficient file not found: {filepath}\n"
@@ -330,6 +336,13 @@ def _load_coefficients_cached(
     actual_n_max = n_max if n_max is not None else int(params["n_max_full"])
     C, S = parse_egm_file(filepath, actual_n_max)
 
+    _logger.info(
+        "Loaded %s coefficients: n_max=%d, array_size=%.1f MB",
+        model,
+        C.shape[0] - 1,
+        C.nbytes / 1024 / 1024 * 2,  # Both C and S arrays
+    )
+
     return EGMCoefficients(
         C=C,
         S=S,

pytcl/gravity/spherical_harmonics.py

@@ -11,11 +11,69 @@ References
 .. [2] O. Montenbruck and E. Gill, "Satellite Orbits," Springer, 2000.
 """
 
+import logging
+from functools import lru_cache
 from typing import Optional, Tuple
 
 import numpy as np
 from numpy.typing import NDArray
 
+# Module logger
+_logger = logging.getLogger("pytcl.gravity.spherical_harmonics")
+
+# Cache configuration for Legendre polynomials
+_LEGENDRE_CACHE_DECIMALS = 8  # Precision for x quantization
+_LEGENDRE_CACHE_MAXSIZE = 64  # Max cached (n_max, m_max, x) combinations
+
+
+def _quantize_x(x: float) -> float:
+    """Quantize x value for cache key compatibility."""
+    return round(x, _LEGENDRE_CACHE_DECIMALS)
+
+
+@lru_cache(maxsize=_LEGENDRE_CACHE_MAXSIZE)
+def _associated_legendre_cached(
+    n_max: int,
+    m_max: int,
+    x_quantized: float,
+    normalized: bool,
+) -> tuple:
+    """Cached Legendre polynomial computation (internal).
+
+    Returns tuple of tuples for hashability.
+    """
+    P = np.zeros((n_max + 1, m_max + 1))
+    u = np.sqrt(1 - x_quantized * x_quantized)
+
+    P[0, 0] = 1.0
+
+    for m in range(1, m_max + 1):
+        if normalized:
+            P[m, m] = u * np.sqrt((2 * m + 1) / (2 * m)) * P[m - 1, m - 1]
+        else:
+            P[m, m] = (2 * m - 1) * u * P[m - 1, m - 1]
+
+    for m in range(m_max):
+        if m + 1 <= n_max:
+            if normalized:
+                P[m + 1, m] = x_quantized * np.sqrt(2 * m + 3) * P[m, m]
+            else:
+                P[m + 1, m] = x_quantized * (2 * m + 1) * P[m, m]
+
+    for m in range(m_max + 1):
+        for n in range(m + 2, n_max + 1):
+            if normalized:
+                a_nm = np.sqrt((4 * n * n - 1) / (n * n - m * m))
+                b_nm = np.sqrt(((n - 1) ** 2 - m * m) / (4 * (n - 1) ** 2 - 1))
+                P[n, m] = a_nm * (x_quantized * P[n - 1, m] - b_nm * P[n - 2, m])
+            else:
+                P[n, m] = (
+                    (2 * n - 1) * x_quantized * P[n - 1, m] - (n + m - 1) * P[n - 2, m]
+                ) / (n - m)
+
+    # Convert to tuple of tuples for hashability
+    return tuple(tuple(row) for row in P)
+
 
 def associated_legendre(
     n_max: int,
@@ -53,6 +111,9 @@ def associated_legendre(
 
         \\int_{-1}^{1} [\\bar{P}_n^m(x)]^2 dx = \\frac{2}{2n+1}
 
+    Results are cached for repeated queries with the same parameters.
+    Cache key quantizes x to 8 decimal places (~1e-8 precision).
+
     Examples
     --------
     >>> P = associated_legendre(2, 2, 0.5)
@@ -63,42 +124,10 @@ def associated_legendre(
     if not -1 <= x <= 1:
         raise ValueError("x must be in [-1, 1]")
 
-    P = np.zeros((n_max + 1, m_max + 1))
-
-    # Compute sqrt(1 - x^2) = sin(theta) for colatitude
-    u = np.sqrt(1 - x * x)
-
-    # Seed values
-    P[0, 0] = 1.0
-
-    # Sectoral recursion: P_m^m from P_{m-1}^{m-1}
-    for m in range(1, m_max + 1):
-        if normalized:
-            P[m, m] = u * np.sqrt((2 * m + 1) / (2 * m)) * P[m - 1, m - 1]
-        else:
-            P[m, m] = (2 * m - 1) * u * P[m - 1, m - 1]
-
-    # Compute P_{m+1}^m from P_m^m
-    for m in range(m_max):
-        if m + 1 <= n_max:
-            if normalized:
-                P[m + 1, m] = x * np.sqrt(2 * m + 3) * P[m, m]
-            else:
-                P[m + 1, m] = x * (2 * m + 1) * P[m, m]
-
-    # General recursion: P_n^m from P_{n-1}^m and P_{n-2}^m
-    for m in range(m_max + 1):
-        for n in range(m + 2, n_max + 1):
-            if normalized:
-                a_nm = np.sqrt((4 * n * n - 1) / (n * n - m * m))
-                b_nm = np.sqrt(((n - 1) ** 2 - m * m) / (4 * (n - 1) ** 2 - 1))
-                P[n, m] = a_nm * (x * P[n - 1, m] - b_nm * P[n - 2, m])
-            else:
-                P[n, m] = (
-                    (2 * n - 1) * x * P[n - 1, m] - (n + m - 1) * P[n - 2, m]
-                ) / (n - m)
-
-    return P
+    # Use cached computation
+    x_q = _quantize_x(x)
+    cached = _associated_legendre_cached(n_max, m_max, x_q, normalized)
+    return np.array(cached)
 
 
 def associated_legendre_derivative(
@@ -230,6 +259,14 @@ def spherical_harmonic_sum(
     if n_max is None:
         n_max = C.shape[0] - 1
 
+    _logger.debug(
+        "spherical_harmonic_sum: lat=%.4f, lon=%.4f, r=%.1f, n_max=%d",
+        lat,
+        lon,
+        r,
+        n_max,
+    )
+
     # Colatitude for Legendre polynomials
     colat = np.pi / 2 - lat
     cos_colat = np.cos(colat)
@@ -495,6 +532,28 @@ def associated_legendre_scaled(
     return P_scaled, scale_exp
 
 
+def clear_legendre_cache() -> None:
+    """Clear cached Legendre polynomial results.
+
+    Call this function to clear the cached associated Legendre
+    polynomial arrays. Useful when memory is constrained or after
+    processing a batch with different colatitude values.
+    """
+    _associated_legendre_cached.cache_clear()
+    _logger.debug("Legendre polynomial cache cleared")
+
+
+def get_legendre_cache_info():
+    """Get cache statistics for Legendre polynomials.
+
+    Returns
+    -------
+    CacheInfo
+        Named tuple with hits, misses, maxsize, currsize.
+    """
+    return _associated_legendre_cached.cache_info()
+
+
 __all__ = [
     "associated_legendre",
     "associated_legendre_derivative",
@@ -502,4 +561,6 @@ __all__ = [
     "gravity_acceleration",
     "legendre_scaling_factors",
     "associated_legendre_scaled",
+    "clear_legendre_cache",
+    "get_legendre_cache_info",
 ]

pytcl/magnetism/__init__.py

@@ -45,7 +45,10 @@ from pytcl.magnetism.wmm import (
     WMM2020,
     MagneticCoefficients,
     MagneticResult,
+    clear_magnetic_cache,
+    configure_magnetic_cache,
     create_wmm2020_coefficients,
+    get_magnetic_cache_info,
     magnetic_declination,
     magnetic_field_intensity,
     magnetic_field_spherical,
@@ -66,6 +69,10 @@ __all__ = [
     "magnetic_declination",
     "magnetic_inclination",
     "magnetic_field_intensity",
+    # Cache management
+    "get_magnetic_cache_info",
+    "clear_magnetic_cache",
+    "configure_magnetic_cache",
     # IGRF
     "IGRF13",
     "create_igrf13_coefficients",

pytcl/magnetism/wmm.py

@@ -12,13 +12,30 @@ References
 .. [2] https://www.ngdc.noaa.gov/geomag/WMM/
 """
 
-from typing import NamedTuple, Tuple
+from functools import lru_cache
+from typing import NamedTuple, Optional, Tuple
 
 import numpy as np
 from numpy.typing import NDArray
 
 from pytcl.gravity.spherical_harmonics import associated_legendre
 
+# =============================================================================
+# Cache Configuration
+# =============================================================================
+
+# Default cache size (number of unique location/time combinations to cache)
+_DEFAULT_CACHE_SIZE = 1024
+
+# Precision for rounding inputs (radians for lat/lon, km for radius, years for time)
+# These control how aggressively similar inputs are grouped
+_CACHE_PRECISION = {
+    "lat": 6,  # ~0.1 meter precision at Earth surface
+    "lon": 6,
+    "r": 3,  # 1 meter precision
+    "year": 2,  # ~4 day precision
+}
+
 
 class MagneticResult(NamedTuple):
     """Result of magnetic field computation.
@@ -404,37 +421,90 @@ def create_wmm2020_coefficients() -> MagneticCoefficients:
 WMM2020 = create_wmm2020_coefficients()
 
 
-def magnetic_field_spherical(
+# =============================================================================
+# Cached Computation Core
+# =============================================================================
+
+
+def _quantize_inputs(
+    lat: float, lon: float, r: float, year: float
+) -> Tuple[float, float, float, float]:
+    """Round inputs to cache precision for consistent cache hits."""
+    return (
+        round(lat, _CACHE_PRECISION["lat"]),
+        round(lon, _CACHE_PRECISION["lon"]),
+        round(r, _CACHE_PRECISION["r"]),
+        round(year, _CACHE_PRECISION["year"]),
+    )
+
+
+@lru_cache(maxsize=_DEFAULT_CACHE_SIZE)
+def _magnetic_field_spherical_cached(
     lat: float,
     lon: float,
     r: float,
     year: float,
-    coeffs: MagneticCoefficients = WMM2020,
+    n_max: int,
+    coeff_id: int,
 ) -> Tuple[float, float, float]:
     """
-    Compute magnetic field in spherical coordinates.
+    Cached core computation of magnetic field in spherical coordinates.
+
+    This is the internal cached version. The coefficient arrays are identified
+    by their id() since NamedTuples with numpy arrays aren't hashable.
 
     Parameters
     ----------
     lat : float
-        Geocentric latitude in radians.
+        Geocentric latitude in radians (quantized).
     lon : float
-        Longitude in radians.
+        Longitude in radians (quantized).
     r : float
-        Radial distance from Earth's center in km.
+        Radial distance in km (quantized).
     year : float
-        Decimal year (e.g., 2023.5 for mid-2023).
-    coeffs : MagneticCoefficients, optional
-        Model coefficients. Default WMM2020.
+        Decimal year (quantized).
+    n_max : int
+        Maximum spherical harmonic degree.
+    coeff_id : int
+        Unique identifier for the coefficient set.
 
     Returns
     -------
-    B_r : float
-        Radial component (positive outward) in nT.
-    B_theta : float
-        Colatitude component (positive southward) in nT.
-    B_phi : float
-        Longitude component (positive eastward) in nT.
+    B_r, B_theta, B_phi : tuple of float
+        Magnetic field components in spherical coordinates (nT).
+    """
+    # Retrieve coefficients from registry
+    coeffs = _coefficient_registry.get(coeff_id)
+    if coeffs is None:
+        raise ValueError(f"Coefficient set {coeff_id} not found in registry")
+
+    return _compute_magnetic_field_spherical_impl(lat, lon, r, year, coeffs)
+
+
+# Registry to hold coefficient sets by id
+_coefficient_registry: dict = {}
+
+
+def _register_coefficients(coeffs: "MagneticCoefficients") -> int:
+    """Register a coefficient set and return its unique ID."""
+    coeff_id = id(coeffs)
+    if coeff_id not in _coefficient_registry:
+        _coefficient_registry[coeff_id] = coeffs
+    return coeff_id
+
+
+def _compute_magnetic_field_spherical_impl(
+    lat: float,
+    lon: float,
+    r: float,
+    year: float,
+    coeffs: "MagneticCoefficients",
+) -> Tuple[float, float, float]:
+    """
+    Core implementation of magnetic field computation.
+
+    This contains the actual spherical harmonic expansion logic,
+    separated for clarity and to support caching.
     """
     n_max = coeffs.n_max
     a = 6371.2  # Reference radius in km (WMM convention)
@@ -452,11 +522,9 @@ def magnetic_field_spherical(
     sin_theta = np.sin(theta)
 
     # Compute associated Legendre functions (Schmidt semi-normalized)
-    # WMM uses Schmidt semi-normalization
     P = associated_legendre(n_max, n_max, cos_theta, normalized=True)
 
-    # Compute dP/dtheta = -sin(theta) * dP/d(cos(theta))
-    # For efficiency, use recurrence relation
+    # Compute dP/dtheta
     dP = np.zeros((n_max + 1, n_max + 1))
     if abs(sin_theta) > 1e-10:
         for n in range(1, n_max + 1):
@@ -464,7 +532,6 @@ def magnetic_field_spherical(
                 if m == n:
                     dP[n, m] = n * cos_theta / sin_theta * P[n, m]
                 elif n > m:
-                    # Recurrence relation for derivative
                     factor = np.sqrt((n - m) * (n + m + 1))
                     if m + 1 <= n:
                         dP[n, m] = (
@@ -489,13 +556,10 @@ def magnetic_field_spherical(
             cos_m_lon = np.cos(m * lon)
             sin_m_lon = np.sin(m * lon)
 
-            # Gauss coefficients
             gnm = g[n, m]
             hnm = h[n, m]
 
-            # Field contributions
             B_r += (n + 1) * r_power * P[n, m] * (gnm * cos_m_lon + hnm * sin_m_lon)
-
             B_theta += -r_power * dP[n, m] * (gnm * cos_m_lon + hnm * sin_m_lon)
 
             if abs(sin_theta) > 1e-10:
@@ -510,6 +574,175 @@ def magnetic_field_spherical(
     return B_r, B_theta, B_phi
 
 
+# =============================================================================
+# Cache Management
+# =============================================================================
+
+
+def get_magnetic_cache_info() -> dict:
+    """
+    Get information about the magnetic field computation cache.
+
+    Returns
+    -------
+    info : dict
+        Dictionary containing cache statistics:
+        - hits: Number of cache hits
+        - misses: Number of cache misses
+        - maxsize: Maximum cache size
+        - currsize: Current number of cached entries
+        - hit_rate: Ratio of hits to total calls (0-1)
+
+    Examples
+    --------
+    >>> from pytcl.magnetism import get_magnetic_cache_info
+    >>> info = get_magnetic_cache_info()
+    >>> print(f"Cache hit rate: {info['hit_rate']:.1%}")
+    """
+    cache_info = _magnetic_field_spherical_cached.cache_info()
+    total = cache_info.hits + cache_info.misses
+    hit_rate = cache_info.hits / total if total > 0 else 0.0
+
+    return {
+        "hits": cache_info.hits,
+        "misses": cache_info.misses,
+        "maxsize": cache_info.maxsize,
+        "currsize": cache_info.currsize,
+        "hit_rate": hit_rate,
+    }
+
+
+def clear_magnetic_cache() -> None:
+    """
+    Clear the magnetic field computation cache.
+
+    This can be useful when memory is constrained or when switching
+    between different coefficient sets.
+
+    Examples
+    --------
+    >>> from pytcl.magnetism import clear_magnetic_cache
+    >>> clear_magnetic_cache()  # Free cached computations
+    """
+    _magnetic_field_spherical_cached.cache_clear()
+    _coefficient_registry.clear()
+
+
+def configure_magnetic_cache(
+    maxsize: Optional[int] = None,
+    precision: Optional[dict] = None,
+) -> None:
+    """
+    Configure the magnetic field computation cache.
+
+    Parameters
+    ----------
+    maxsize : int, optional
+        Maximum number of entries in the cache. If None, keeps current.
+        Set to 0 to disable caching.
+    precision : dict, optional
+        Dictionary with keys 'lat', 'lon', 'r', 'year' specifying
+        decimal places for rounding. Higher values = more precision
+        but fewer cache hits.
+
+    Notes
+    -----
+    Changing cache configuration clears the existing cache.
+
+    Examples
+    --------
+    >>> from pytcl.magnetism import configure_magnetic_cache
+    >>> # Increase cache size for batch processing
+    >>> configure_magnetic_cache(maxsize=4096)
+    >>> # Reduce precision for more cache hits
+    >>> configure_magnetic_cache(precision={'lat': 4, 'lon': 4, 'r': 2, 'year': 1})
+    """
+    global _magnetic_field_spherical_cached
+
+    if precision is not None:
+        for key in ["lat", "lon", "r", "year"]:
+            if key in precision:
+                _CACHE_PRECISION[key] = precision[key]
+
+    if maxsize is not None:
+        # Recreate the cached function with new maxsize
+        clear_magnetic_cache()
+
+        @lru_cache(maxsize=maxsize)
+        def new_cached(
+            lat: float,
+            lon: float,
+            r: float,
+            year: float,
+            n_max: int,
+            coeff_id: int,
+        ) -> Tuple[float, float, float]:
+            coeffs = _coefficient_registry.get(coeff_id)
+            if coeffs is None:
+                raise ValueError(f"Coefficient set {coeff_id} not found")
+            return _compute_magnetic_field_spherical_impl(lat, lon, r, year, coeffs)
+
+        _magnetic_field_spherical_cached = new_cached
+
+
+def magnetic_field_spherical(
+    lat: float,
+    lon: float,
+    r: float,
+    year: float,
+    coeffs: MagneticCoefficients = WMM2020,
+    use_cache: bool = True,
+) -> Tuple[float, float, float]:
+    """
+    Compute magnetic field in spherical coordinates.
+
+    Parameters
+    ----------
+    lat : float
+        Geocentric latitude in radians.
+    lon : float
+        Longitude in radians.
+    r : float
+        Radial distance from Earth's center in km.
+    year : float
+        Decimal year (e.g., 2023.5 for mid-2023).
+    coeffs : MagneticCoefficients, optional
+        Model coefficients. Default WMM2020.
+    use_cache : bool, optional
+        Whether to use LRU caching for repeated queries. Default True.
+        Set to False for single-use queries or when memory is constrained.
+
+    Returns
+    -------
+    B_r : float
+        Radial component (positive outward) in nT.
+    B_theta : float
+        Colatitude component (positive southward) in nT.
+    B_phi : float
+        Longitude component (positive eastward) in nT.
+
+    Notes
+    -----
+    Results are cached by default using LRU caching. Inputs are quantized
+    to a configurable precision before caching to improve hit rates for
+    nearby queries. Use `get_magnetic_cache_info()` to check cache
+    statistics and `clear_magnetic_cache()` to free memory.
+    """
+    if use_cache:
+        # Quantize inputs for cache key
+        q_lat, q_lon, q_r, q_year = _quantize_inputs(lat, lon, r, year)
+
+        # Register coefficients and get ID
+        coeff_id = _register_coefficients(coeffs)
+
+        return _magnetic_field_spherical_cached(
+            q_lat, q_lon, q_r, q_year, coeffs.n_max, coeff_id
+        )
+    else:
+        # Direct computation without caching
+        return _compute_magnetic_field_spherical_impl(lat, lon, r, year, coeffs)
+
+
 def wmm(
     lat: float,
     lon: float,
@@ -706,4 +939,8 @@ __all__ = [
     "magnetic_declination",
     "magnetic_inclination",
     "magnetic_field_intensity",
+    # Cache management
+    "get_magnetic_cache_info",
+    "clear_magnetic_cache",
+    "configure_magnetic_cache",
 ]