nlsq.memory_manager module¶

Memory management for NLSQ optimization.

This module provides intelligent memory management capabilities including prediction, monitoring, pooling, and automatic garbage collection.

Phase 3 Optimizations (Task Group 9):

Telemetry Circular Buffer (1.3a): Uses deque(maxlen=1000) for _safety_telemetry to prevent memory leak in multi-day optimization runs

class nlsq.caching.memory_manager.MemoryManager(gc_threshold=0.8, safety_factor=1.2, enable_adaptive_safety=False, disable_padding=False, memory_cache_ttl=1.0, adaptive_ttl=True)[source]

Bases: object

Intelligent memory management for optimization algorithms.

This class provides: - Memory usage monitoring and prediction - Array pooling to reduce allocations with LRU eviction - Automatic garbage collection triggers - Context managers for memory-safe operations

LRU Memory Pool (Task Group 7 - 1.2a)¶

The memory pool uses an OrderedDict to track access order, enabling true LRU (Least Recently Used) eviction when at capacity. This improves cache utilization for frequently accessed array shapes by 5-10%.

Telemetry Circular Buffer (Task Group 9 - 1.3a)¶

The safety telemetry uses a deque with maxlen=1000 to prevent memory leak in multi-day optimization runs. This maintains the last 1000 telemetry records for adaptive safety factor calculation.

memory_pool

Pool of reusable arrays indexed by (shape, dtype) with LRU tracking

Type:: OrderedDict

allocation_history

History of memory allocations

Type:: list

gc_threshold

Memory usage threshold (0-1) for triggering garbage collection

Type:: float

safety_factor

Safety factor for memory predictions

Type:: float

__init__(gc_threshold=0.8, safety_factor=1.2, enable_adaptive_safety=False, disable_padding=False, memory_cache_ttl=1.0, adaptive_ttl=True)[source]

Initialize memory manager.

Parameters:

gc_threshold (float) – Trigger GC when memory usage exceeds this fraction (0-1)
safety_factor (float) – Multiply memory requirements by this factor for safety
enable_adaptive_safety (bool) – Enable adaptive safety factor reduction (1.2 -> 1.05 after warmup)
disable_padding (bool) – Disable padding/bucketing for strict memory environments (Task 5.6). When True: uses exact shapes, sets safety_factor=1.0. Use case: cloud quotas, strict memory limits.
memory_cache_ttl (float) – TTL in seconds for cached memory info (default: 1.0). Reduces psutil system call overhead by 90%.
adaptive_ttl (bool) – Enable adaptive TTL based on call frequency (default: True). High-frequency callers (>100 calls/sec) get 15s effective TTL. Medium-frequency callers (>10 calls/sec) get 10s effective TTL. Low-frequency callers use the default TTL. Reduces psutil overhead in streaming optimization by 15-20%.

get_available_memory()[source]

Get available memory in bytes.

Returns:: available – Available memory in bytes
Return type:: float

Notes

Uses TTL-based caching to reduce psutil system call overhead by 90%. When adaptive_ttl is enabled, the effective TTL is adjusted based on call frequency to further reduce overhead for streaming optimization.

get_memory_usage_bytes()[source]

Get current memory usage in bytes.

Returns:: usage – Current memory usage in bytes
Return type:: float

Notes

Uses TTL-based caching to reduce psutil system call overhead by 90%.

get_memory_usage_fraction()[source]

Get current memory usage as fraction of total.

Returns:: fraction – Memory usage fraction (0-1)
Return type:: float

Notes

Uses TTL-based caching to reduce psutil system call overhead by 90%.

predict_memory_requirement(n_points, n_params, algorithm='trf', dtype=<class 'jax.numpy.float64'>)[source]

Predict memory requirement for optimization.

Parameters:

n_points (int) – Number of data points
n_params (int) – Number of parameters
algorithm (str) – Algorithm name (‘trf’, ‘lm’, ‘dogbox’)
dtype (jnp.dtype, optional) – Data type for computations (default: jnp.float64). Affects memory calculations: float32 uses 4 bytes, float64 uses 8 bytes.

Returns:

bytes_needed – Estimated memory requirement in bytes

Return type:

int

Notes

Memory requirements scale linearly with precision: - float32: 4 bytes per element (50% memory savings) - float64: 8 bytes per element (default, higher precision)

check_memory_availability(bytes_needed)[source]

Check if enough memory is available.

Parameters:

bytes_needed (int) – Memory required in bytes

Returns:

available (bool) – Whether enough memory is available
message (str) – Descriptive message

Return type:

tuple[bool, str]

memory_guard(bytes_needed)[source]

Context manager to ensure memory availability.

Parameters:: bytes_needed (int) – Required memory in bytes
Raises:: MemoryError – If insufficient memory is available

get_safety_telemetry()[source]

Get safety factor telemetry statistics.

Returns:: telemetry – Safety factor telemetry with: - current_safety_factor: Current safety factor - initial_safety_factor: Initial safety factor (1.2) - min_safety_factor: Target minimum (1.05) - telemetry_entries: Number of telemetry entries collected - p95_safety_needed: 95th percentile of safety factors needed (if data available) - safety_factor_history: List of safety factors over time
Return type:: dict

allocate_array(shape, dtype=<class 'numpy.float64'>, zero=True)[source]

Allocate array with memory pooling and LRU tracking.

Parameters:

shape (tuple) – Shape of array to allocate
dtype (type) – Data type of array
zero (bool) – Whether to zero-initialize the array

Returns:

array – Allocated array

Return type:

np.ndarray

Raises:

MemoryError – If allocation fails

Notes

Task Group 7 (1.2a): Uses LRU tracking via OrderedDict. When an array is reused from the pool, it is moved to the end (most recently used) to enable proper LRU eviction.

free_array(arr)[source]

Return array to pool for reuse.

Parameters:: arr (np.ndarray) – Array to free

Notes

Task Group 7 (1.2a): Uses LRU tracking via OrderedDict. The returned array is added/moved to the end of the pool, marking it as recently used.

clear_pool()[source]: Clear memory pool and run garbage collection.

get_memory_stats()[source]

Get memory usage statistics.

Returns:: stats – Memory statistics including current usage, peak, pool size
Return type:: dict

optimize_memory_pool(max_arrays=100)[source]

Optimize memory pool using LRU eviction.

Parameters:: max_arrays (int) – Maximum number of arrays to keep in pool

Notes

Task Group 7 (1.2a): Uses LRU eviction via popitem(last=False). Arrays are evicted in order of least recent use, keeping the most recently used arrays in the pool.

temporary_allocation(shape, dtype=<class 'numpy.float64'>)[source]

Context manager for temporary array allocation.

Parameters:

shape (tuple) – Shape of array
dtype (type) – Data type

Yields:

array (np.ndarray) – Temporary array that will be returned to pool on exit

estimate_chunking_strategy(n_points, n_params, algorithm='trf', memory_limit_gb=None)[source]

Estimate optimal chunking strategy for large datasets.

Parameters:

n_points (int) – Number of data points
n_params (int) – Number of parameters
algorithm (str) – Algorithm to use
memory_limit_gb (float, optional) – Memory limit in GB (uses available memory if None)

Returns:

strategy – Chunking strategy with chunk_size and n_chunks

Return type:

dict

nlsq.caching.memory_manager.get_memory_manager()[source]

Get or create global memory manager instance.

Returns:: manager – Global memory manager instance
Return type:: MemoryManager

nlsq.caching.memory_manager.clear_memory_pool()[source]: Clear the global memory pool.

nlsq.caching.memory_manager.get_memory_stats()[source]

Get memory usage statistics.

Returns:: stats – Memory statistics
Return type:: dict

Overview¶

The nlsq.memory_manager module provides intelligent memory management capabilities for optimization algorithms. It includes memory usage monitoring, prediction, array pooling, and automatic garbage collection to handle large-scale curve fitting problems efficiently.

Key Features¶

Memory usage monitoring with psutil integration
Memory requirement prediction for different algorithms
Array pooling to reduce allocation overhead
Automatic garbage collection triggers
Memory-safe context managers for operations
Chunking strategy estimation for large datasets
Allocation history tracking and statistics

Classes¶

MemoryManager([gc_threshold, safety_factor, ...])

Intelligent memory management for optimization algorithms.

Functions¶

`get_memory_manager`()	Get or create global memory manager instance.
`clear_memory_pool`()	Clear the global memory pool.
`get_memory_stats`()	Get memory usage statistics.

Usage Examples¶

Basic Memory Management¶

Monitor and manage memory during optimization:

from nlsq.caching.memory_manager import MemoryManager
import numpy as np

# Create memory manager
mm = MemoryManager(gc_threshold=0.8, safety_factor=1.2)

# Check available memory
available = mm.get_available_memory()
print(f"Available memory: {available / 1e9:.2f} GB")

# Predict memory requirements
n_points = 1_000_000
n_params = 10
bytes_needed = mm.predict_memory_requirement(n_points, n_params, algorithm="trf")
print(f"Memory needed: {bytes_needed / 1e9:.2f} GB")

# Check if enough memory is available
is_available, message = mm.check_memory_availability(bytes_needed)
if is_available:
    print("Sufficient memory available")
else:
    print(f"Warning: {message}")

Memory-Safe Operations¶

Use context managers to ensure memory availability:

from nlsq.caching.memory_manager import MemoryManager
import numpy as np

mm = MemoryManager()

# Estimate memory needed
n_points = 5_000_000
n_params = 5
bytes_needed = mm.predict_memory_requirement(n_points, n_params, "trf")

# Use memory guard to ensure availability
try:
    with mm.memory_guard(bytes_needed):
        # Perform memory-intensive operation
        x = np.random.randn(n_points)
        y = 2 * x + 1 + np.random.randn(n_points) * 0.1

        # Your optimization here
        result = optimize(x, y)

except MemoryError as e:
    print(f"Insufficient memory: {e}")
    # Fall back to chunked processing

Array Pooling¶

Reuse arrays to reduce allocation overhead:

from nlsq.caching.memory_manager import MemoryManager
import numpy as np

mm = MemoryManager()

# Allocate arrays from pool
jacobian = mm.allocate_array((1000, 10), dtype=np.float64, zero=True)
residuals = mm.allocate_array((1000,), dtype=np.float64, zero=True)

# Use arrays
jacobian[:] = compute_jacobian()
residuals[:] = compute_residuals()

# Return arrays to pool when done
mm.free_array(jacobian)
mm.free_array(residuals)

# Arrays will be reused on next allocation
jacobian2 = mm.allocate_array((1000, 10), dtype=np.float64)  # Reuses pooled array

Temporary Allocations¶

Use context manager for temporary arrays:

from nlsq.caching.memory_manager import MemoryManager

mm = MemoryManager()

# Temporary array automatically returned to pool
with mm.temporary_allocation((1000, 50), dtype=np.float64) as temp_array:
    # Use temporary array
    temp_array[:] = compute_intermediate_result()
    final_result = process(temp_array)

# temp_array is now back in the pool for reuse

Chunking Strategy Estimation¶

Estimate optimal chunk sizes for large datasets:

from nlsq.caching.memory_manager import MemoryManager

mm = MemoryManager()

# Estimate chunking for 100M points
n_points = 100_000_000
n_params = 5
memory_limit_gb = 4.0

strategy = mm.estimate_chunking_strategy(
    n_points, n_params, algorithm="trf", memory_limit_gb=memory_limit_gb
)

if strategy["needs_chunking"]:
    print(f"Chunking required:")
    print(f"  Chunk size: {strategy['chunk_size']:,} points")
    print(f"  Number of chunks: {strategy['n_chunks']}")
    print(f"  Memory per chunk: {strategy['memory_per_chunk_gb']:.2f} GB")
else:
    print("No chunking needed - dataset fits in memory")

Global Memory Manager¶

Use the global memory manager instance:

from nlsq.caching.memory_manager import (
    get_memory_manager,
    get_memory_stats,
    clear_memory_pool,
)

# Get global instance
mm = get_memory_manager()

# Use it
arr = mm.allocate_array((1000, 100))

# Get memory statistics
stats = get_memory_stats()
print(f"Current usage: {stats['current_usage_gb']:.2f} GB")
print(f"Peak usage: {stats['peak_usage_gb']:.2f} GB")
print(f"Pool size: {stats['pool_arrays']} arrays")
print(f"Allocation efficiency: {stats.get('efficiency', 1.0):.1%}")

# Clear pool when done
clear_memory_pool()

Integration with Large Dataset Fitting¶

Memory manager integrates with large dataset tools:

from nlsq import curve_fit_large
from nlsq.caching.memory_manager import MemoryManager
import jax.numpy as jnp
import numpy as np

# Set up memory manager
mm = MemoryManager(gc_threshold=0.75)

# Check memory requirements
n_points = 50_000_000
n_params = 3
bytes_needed = mm.predict_memory_requirement(n_points, n_params, "trf")

is_available, message = mm.check_memory_availability(bytes_needed)

if not is_available:
    # Use chunking
    strategy = mm.estimate_chunking_strategy(n_points, n_params, memory_limit_gb=4.0)
    print(f"Using {strategy['n_chunks']} chunks")


# Fit with automatic memory management
def exponential(x, a, b, c):
    return a * jnp.exp(-b * x) + c


x = np.linspace(0, 10, n_points)
y = 2.0 * np.exp(-0.5 * x) + 0.3 + np.random.normal(0, 0.05, n_points)

with mm.memory_guard(bytes_needed):
    popt, pcov = curve_fit_large(
        exponential, x, y, p0=[2.5, 0.6, 0.2], memory_limit_gb=4.0
    )

Memory Statistics and Diagnostics¶

Track memory usage over time:

from nlsq.caching.memory_manager import MemoryManager

mm = MemoryManager()

# Perform operations
for i in range(10):
    arr = mm.allocate_array((10000, 100))
    # ... use array ...
    mm.free_array(arr)

# Get detailed statistics
stats = mm.get_memory_stats()

print("Memory Statistics:")
print(f"  Current usage: {stats['current_usage_gb']:.2f} GB")
print(f"  Available: {stats['available_gb']:.2f} GB")
print(f"  Peak usage: {stats['peak_usage_gb']:.2f} GB")
print(f"  Usage fraction: {stats['usage_fraction']:.1%}")
print(f"  Pool memory: {stats['pool_memory_gb']:.3f} GB")
print(f"  Pool arrays: {stats['pool_arrays']}")
print(f"  Total allocations: {stats['allocations']}")

if "efficiency" in stats:
    print(f"  Allocation efficiency: {stats['efficiency']:.1%}")
    print(f"  Total requested: {stats['total_requested_gb']:.2f} GB")
    print(f"  Total used: {stats['total_used_gb']:.2f} GB")

# Optimize pool if it grows too large
mm.optimize_memory_pool(max_arrays=50)

Memory Prediction for Different Algorithms¶

Compare memory requirements across algorithms:

from nlsq.caching.memory_manager import MemoryManager

mm = MemoryManager()

n_points = 10_000
n_params = 20

algorithms = ["trf", "lm", "dogbox"]
for algo in algorithms:
    bytes_needed = mm.predict_memory_requirement(n_points, n_params, algo)
    print(f"{algo:8s}: {bytes_needed / 1e6:.2f} MB")

# Output:
# trf     : 12.34 MB
# lm      : 8.76 MB
# dogbox  : 13.45 MB

Configuration¶

Memory manager can be configured at initialization:

from nlsq.caching.memory_manager import MemoryManager

# Conservative settings (more GC, larger safety margin)
conservative_mm = MemoryManager(gc_threshold=0.7, safety_factor=1.5)

# Aggressive settings (less GC, smaller safety margin)
aggressive_mm = MemoryManager(gc_threshold=0.9, safety_factor=1.1)

# Custom settings for specific use case
custom_mm = MemoryManager(
    gc_threshold=0.8,  # Trigger GC at 80% memory usage
    safety_factor=1.3,  # Add 30% safety margin to predictions
)

Performance Considerations¶

Array pooling benefits:

Reduces allocation overhead by 50-80% for repeated operations
Most beneficial for medium-sized arrays (10 KB - 10 MB)
Pool automatically optimized to keep largest arrays

Memory prediction accuracy:

Predictions are conservative (safety_factor=1.2 by default)
Accuracy: ±10-20% for standard algorithms
Algorithm-specific formulas based on known memory patterns

Garbage collection:

Triggered automatically when usage exceeds gc_threshold
Manual collection available via clear_pool()
Recommended to clear pool between independent fitting operations

psutil dependency:

Optional but recommended for accurate memory monitoring
Falls back to conservative estimates if unavailable
Install with: pip install psutil

Memory Efficiency Tips¶

Reuse arrays: Use array pooling for repeated allocations
Clear pool periodically: Especially between independent tasks
Estimate before allocating: Check availability with memory_guard()
Use chunking: For datasets that don’t fit in memory
Monitor statistics: Track efficiency with get_memory_stats()
Optimize pool size: Keep pool under 100 arrays for best performance

nlsq.memory_manager module¶

LRU Memory Pool (Task Group 7 - 1.2a)¶

Telemetry Circular Buffer (Task Group 9 - 1.3a)¶

Overview¶

Key Features¶

Classes¶

Functions¶

Usage Examples¶

Basic Memory Management¶

Memory-Safe Operations¶

Array Pooling¶

Temporary Allocations¶

Chunking Strategy Estimation¶

Global Memory Manager¶

Integration with Large Dataset Fitting¶

Memory Statistics and Diagnostics¶

Memory Prediction for Different Algorithms¶

Configuration¶

Performance Considerations¶

Memory Efficiency Tips¶

See Also¶