nlsq.smart_cache module

Smart caching system for NLSQ optimization.

This module provides intelligent caching for expensive computations, particularly Jacobian evaluations and function calls.

Note: This module uses safe serialization only (JSON and numpy.savez with allow_pickle=False). No pickle is used.

Phase 3 Optimizations (Task Group 9): - Array hash optimization: stride-based sampling only for >10000 elements - For smaller arrays, hash full array directly (no redundant sampling)

class nlsq.caching.smart_cache.SmartCache(cache_dir='.nlsq_cache', max_memory_items=1000, disk_cache_enabled=True, enable_stats=True)[source]

Bases: object

Intelligent caching system for optimization computations.

This class provides:

  • Memory and disk caching with LRU eviction

  • Automatic cache key generation from function arguments

  • Cache persistence across sessions

  • Cache invalidation and warming strategies

Phase 3 Optimizations (3.2a):

  • Array hash optimization: uses stride-based sampling only for arrays with >10000 elements when xxhash is unavailable

  • For smaller arrays, hashes full array directly without redundant sampling, providing 15-20% improvement in cache key generation

All dict operations are protected by a per-instance threading.Lock so that concurrent threads can safely call get/set/invalidate.

cache_dir

Directory for disk cache storage

Type:

str

memory_cache

In-memory cache storage

Type:

dict

disk_cache_enabled

Whether disk caching is enabled

Type:

bool

max_memory_items

Maximum items in memory cache

Type:

int

cache_stats

Cache hit/miss statistics

Type:

dict

__init__(cache_dir='.nlsq_cache', max_memory_items=1000, disk_cache_enabled=True, enable_stats=True)[source]

Initialize smart cache.

Parameters:

  • cache_dir (str) – Directory for disk cache

  • max_memory_items (int) – Maximum items in memory cache

  • disk_cache_enabled (bool) – Enable disk caching

  • enable_stats (bool) – Track cache statistics

cache_key(*args, **kwargs)[source]

Generate cache key from arguments.

Parameters:

  • *args (tuple) – Positional arguments

  • **kwargs (dict) – Keyword arguments

Returns:

key – Hash of arguments (xxhash if available, BLAKE2b fallback)

Return type:

str

Notes

Uses xxhash (xxh64) when available for ~10x faster hashing compared to SHA256/BLAKE2b. Falls back to BLAKE2b if xxhash is not installed. All cache keys are prefixed with CACHE_VERSION to ensure old cache entries are invalidated when the hash algorithm changes.

Task 9.3 (3.2a): Array hash optimization - Arrays <= 10000 elements: hash full array directly (no sampling) - Arrays > 10000 elements: use stride-based sampling for efficiency - Removes redundant sampling when computing full hash in fallback path

get(key)[source]

Get value from cache.

Parameters:

key (str) – Cache key

Returns:

value – Cached value or None if not found

Return type:

Any or None

set(key, value)[source]

Set value in cache.

Parameters:

  • key (str) – Cache key

  • value (Any) – Value to cache

invalidate(key=None)[source]

Invalidate cache entries.

Parameters:

key (str, optional) – Specific key to invalidate, or None to clear all

get_stats()[source]

Get cache statistics.

Returns:

stats – Cache statistics including hit rate

Return type:

dict

optimize_cache()[source]

Optimize cache by removing rarely accessed items.

Computes threshold from snapshot, then re-checks live counts under lock before invalidating to avoid evicting keys that became hot between the snapshot and eviction.

nlsq.caching.smart_cache.cached_function(cache=None, ttl=None)[source]

Decorator for caching function results.

Parameters:

  • cache (SmartCache, optional) – Cache instance to use (creates new if None)

  • ttl (float, optional) – Time-to-live in seconds for cached values

Returns:

decorator – Decorator function

Return type:

function

nlsq.caching.smart_cache.cached_jacobian(cache=None)[source]

Decorator specifically for caching Jacobian evaluations.

Parameters:

cache (SmartCache, optional) – Cache instance to use

Returns:

decorator – Decorator function

Return type:

function

class nlsq.caching.smart_cache.JITCompilationCache(max_cache_size=256)[source]

Bases: object

Cache for JAX JIT-compiled functions with LRU eviction.

This cache stores compiled functions to avoid recompilation when function signatures match. Uses OrderedDict for LRU eviction to prevent unbounded XLA compilation cache growth.

All dict operations are protected by a per-instance threading.Lock so that concurrent threads can safely call get_or_compile/clear.

Parameters:

max_cache_size (int) – Maximum number of compiled functions to cache (default 256). Oldest entries are evicted when capacity is reached.

__init__(max_cache_size=256)[source]

Initialize JIT compilation cache with LRU eviction.

get_or_compile(func, static_argnums=())[source]

Get cached compilation or compile and cache.

Parameters:

  • func (callable) – Function to compile

  • static_argnums (tuple) – Static argument numbers for JIT

Returns:

compiled_func – JIT-compiled function

Return type:

callable

clear()[source]

Clear compilation cache.

get_stats()[source]

Get compilation statistics.

Returns:

stats – Compilation statistics

Return type:

dict

nlsq.caching.smart_cache.get_global_cache()[source]

Get global cache instance.

Returns:

cache – Global cache instance

Return type:

SmartCache

nlsq.caching.smart_cache.get_jit_cache()[source]

Get JIT compilation cache.

Returns:

cache – JIT compilation cache

Return type:

JITCompilationCache

nlsq.caching.smart_cache.clear_all_caches()[source]

Clear all global caches.

Overview

The smart_cache module provides intelligent caching with adaptive features for JIT compilation and function evaluation.

Key Features

  • Adaptive cache sizing based on usage patterns

  • Memory-aware eviction policies with size limits

  • Function evaluation caching for repeated calls

  • Jacobian caching with automatic invalidation

  • Statistics tracking for cache hit rate monitoring

Classes

class nlsq.caching.smart_cache.SmartCache(cache_dir='.nlsq_cache', max_memory_items=1000, disk_cache_enabled=True, enable_stats=True)[source]

Bases: object

Intelligent caching system for optimization computations.

This class provides:

  • Memory and disk caching with LRU eviction

  • Automatic cache key generation from function arguments

  • Cache persistence across sessions

  • Cache invalidation and warming strategies

Phase 3 Optimizations (3.2a):

  • Array hash optimization: uses stride-based sampling only for arrays with >10000 elements when xxhash is unavailable

  • For smaller arrays, hashes full array directly without redundant sampling, providing 15-20% improvement in cache key generation

All dict operations are protected by a per-instance threading.Lock so that concurrent threads can safely call get/set/invalidate.

cache_dir

Directory for disk cache storage

Type:

str

memory_cache

In-memory cache storage

Type:

dict

disk_cache_enabled

Whether disk caching is enabled

Type:

bool

max_memory_items

Maximum items in memory cache

Type:

int

cache_stats

Cache hit/miss statistics

Type:

dict

__init__(cache_dir='.nlsq_cache', max_memory_items=1000, disk_cache_enabled=True, enable_stats=True)[source]

Initialize smart cache.

Parameters:

  • cache_dir (str) – Directory for disk cache

  • max_memory_items (int) – Maximum items in memory cache

  • disk_cache_enabled (bool) – Enable disk caching

  • enable_stats (bool) – Track cache statistics

cache_key(*args, **kwargs)[source]

Generate cache key from arguments.

Parameters:

  • *args (tuple) – Positional arguments

  • **kwargs (dict) – Keyword arguments

Returns:

key – Hash of arguments (xxhash if available, BLAKE2b fallback)

Return type:

str

Notes

Uses xxhash (xxh64) when available for ~10x faster hashing compared to SHA256/BLAKE2b. Falls back to BLAKE2b if xxhash is not installed. All cache keys are prefixed with CACHE_VERSION to ensure old cache entries are invalidated when the hash algorithm changes.

Task 9.3 (3.2a): Array hash optimization - Arrays <= 10000 elements: hash full array directly (no sampling) - Arrays > 10000 elements: use stride-based sampling for efficiency - Removes redundant sampling when computing full hash in fallback path

get(key)[source]

Get value from cache.

Parameters:

key (str) – Cache key

Returns:

value – Cached value or None if not found

Return type:

Any or None

set(key, value)[source]

Set value in cache.

Parameters:

  • key (str) – Cache key

  • value (Any) – Value to cache

invalidate(key=None)[source]

Invalidate cache entries.

Parameters:

key (str, optional) – Specific key to invalidate, or None to clear all

get_stats()[source]

Get cache statistics.

Returns:

stats – Cache statistics including hit rate

Return type:

dict

optimize_cache()[source]

Optimize cache by removing rarely accessed items.

Computes threshold from snapshot, then re-checks live counts under lock before invalidating to avoid evicting keys that became hot between the snapshot and eviction.

Functions

nlsq.caching.smart_cache.cached_function(cache=None, ttl=None)[source]

Decorator for caching function results.

Parameters:

  • cache (SmartCache, optional) – Cache instance to use (creates new if None)

  • ttl (float, optional) – Time-to-live in seconds for cached values

Returns:

decorator – Decorator function

Return type:

function

Example Usage

from nlsq.caching.smart_cache import SmartCache, cached_function
from nlsq import curve_fit
import jax.numpy as jnp

# Configure caching
cache = SmartCache(max_memory_items=1000, disk_cache_enabled=True)


# Define fit function
@cached_function(cache=cache)
def exponential(x, a, b):
    return a * jnp.exp(-b * x)


# First fit - compiles function
x1 = jnp.linspace(0, 10, 100)
y1 = 2.5 * jnp.exp(-0.5 * x1) + 0.1 * jnp.random.randn(100)
popt1, pcov1 = curve_fit(exponential, x1, y1, p0=[1.0, 0.1])

# Second fit - reuses JIT compilation from first fit
x2 = jnp.linspace(0, 10, 100)
y2 = 3.0 * jnp.exp(-0.7 * x2) + 0.1 * jnp.random.randn(100)
popt2, pcov2 = curve_fit(exponential, x2, y2, p0=[1.2, 0.15])

# Check cache statistics
stats = cache.get_stats()
print(f"Cache hit rate: {stats['hit_rate']:.1%}")
print(f"Total compilations: {stats['compilations']}")

Cache Management

# Create cache with custom settings
cache = SmartCache(
    max_memory_items=500,
    disk_cache_enabled=True,
    eviction_policy="lru",  # Least Recently Used
    max_disk_size_mb=100,
)

# Manual cache control
cache.clear()  # Clear all cached items
cache.evict_lru()  # Evict least recently used items

# Get detailed statistics
stats = cache.get_detailed_stats()
print(f"Memory usage: {stats['memory_mb']:.2f} MB")
print(f"Disk usage: {stats['disk_mb']:.2f} MB")
print(f"Hit rate: {stats['hit_rate']:.2%}")

See Also