nlsq.large_dataset module

Large Dataset Fitting Module for NLSQ.

This module provides utilities for efficiently fitting curve parameters to very large datasets (>10M points) with intelligent memory management, automatic chunking, and progress reporting.

class nlsq.streaming.large_dataset.ChunkBufferPool(chunk_size, dtype=<class 'numpy.float64'>)[source]

Bases: object

Pre-allocated buffer pool for streaming chunk operations (FR-006).

Reduces memory allocation overhead by reusing buffers across chunks. This is particularly beneficial for large dataset processing where many chunks of similar size are processed sequentially.

Parameters:

  • chunk_size (int) – Size of each chunk buffer

  • dtype (np.dtype) – Data type for buffers (default: np.float64)

Examples

>>> pool = ChunkBufferPool(chunk_size=10000)
>>> x_buf, y_buf = pool.get_buffers(5000)
>>> x_buf.shape
(5000,)
>>> # Buffers can be reused without reallocation
>>> x_buf2, y_buf2 = pool.get_buffers(8000)
>>> x_buf2.shape
(8000,)
__init__(chunk_size, dtype=<class 'numpy.float64'>)[source]

Initialize buffer pool with pre-allocated arrays.

Parameters:

  • chunk_size (int) – Size of each chunk buffer

  • dtype (np.dtype) – Data type for buffers

get_buffers(size)[source]

Get buffer views of the specified size.

Parameters:

size (int) – Actual size needed (must be <= chunk_size)

Returns:

Views into the pre-allocated x and y buffers

Return type:

tuple[np.ndarray, np.ndarray]

property chunk_size: int

Return the pool’s chunk size.

property dtype: dtype

Return the pool’s data type.

class nlsq.streaming.large_dataset.DataChunker[source]

Bases: object

Utility for creating and managing data chunks.

static create_chunks(xdata, ydata, chunk_size, shuffle=False, random_seed=None, buffer_pool=None)[source]

Create data chunks for processing.

Parameters:

  • xdata (np.ndarray) – Independent variable data

  • ydata (np.ndarray) – Dependent variable data

  • chunk_size (int) – Size of each chunk

  • shuffle (bool, optional) – Whether to shuffle data before chunking (default: False)

  • random_seed (int, optional) – Random seed for shuffling

  • buffer_pool (ChunkBufferPool, optional) – Pre-allocated buffer pool for chunk reuse (FR-006). If provided, chunks are written into pooled buffers to reduce allocation overhead.

Yields:

tuple[np.ndarray, np.ndarray, int, int] – (x_chunk, y_chunk, chunk_index, valid_length) where valid_length is the actual number of data points (before padding)

Notes

Uses power-of-2 bucket sizes from CHUNK_BUCKETS for JIT cache efficiency. This ensures consistent array shapes across chunks, enabling JAX to reuse compiled kernels and avoiding recompilation overhead.

class nlsq.streaming.large_dataset.DatasetStats(n_points, n_params, memory_per_point_bytes, total_memory_estimate_gb, recommended_chunk_size, n_chunks)[source]

Bases: object

Statistics and information about a dataset.

n_points

Total number of data points

Type:

int

n_params

Number of parameters to fit

Type:

int

memory_per_point_bytes

Estimated memory usage per data point in bytes

Type:

float

total_memory_estimate_gb

Estimated total memory requirement in GB

Type:

float

recommended_chunk_size

Recommended chunk size for processing

Type:

int

n_chunks

Number of chunks needed

Type:

int

n_points: int
n_params: int
memory_per_point_bytes: float
total_memory_estimate_gb: float
recommended_chunk_size: int
n_chunks: int
__init__(n_points, n_params, memory_per_point_bytes, total_memory_estimate_gb, recommended_chunk_size, n_chunks)
class nlsq.streaming.large_dataset.GPUMemoryEstimator[source]

Bases: object

Utilities for estimating GPU memory availability.

This class provides GPU memory detection via JAX’s device API, handling multiple GPUs and graceful fallback for CPU-only environments.

Examples

>>> from nlsq.streaming.large_dataset import GPUMemoryEstimator
>>> estimator = GPUMemoryEstimator()
>>> available_gb = estimator.get_available_gpu_memory_gb()
>>> print(f"Available GPU memory: {available_gb:.2f} GB")
__init__()[source]

Initialize GPUMemoryEstimator.

get_available_gpu_memory_gb()[source]

Get available GPU memory in GB.

Queries GPU memory via jax.devices()[i].memory_stats() and aggregates available memory across all GPUs. Returns 0 for CPU-only environments.

Returns:

Available GPU memory in GB, or 0.0 if no GPU or detection fails.

Return type:

float

Notes

  • Re-evaluates on each call (no caching) per requirements.

  • Handles multiple GPUs by summing available memory.

  • Returns 0.0 gracefully for CPU-only environments or when detection fails.

get_total_gpu_memory_gb()[source]

Get total GPU memory capacity in GB.

Returns:

Total GPU memory capacity in GB, or 0.0 if no GPU.

Return type:

float

has_gpu()[source]

Check if any GPU is available.

Returns:

True if at least one GPU device is available.

Return type:

bool

class nlsq.streaming.large_dataset.LDMemoryConfig(memory_limit_gb=8.0, safety_factor=0.8, min_chunk_size=1000, max_chunk_size=1000000, use_streaming=True, streaming_batch_size=50000, streaming_max_epochs=10, min_success_rate=0.5, save_diagnostics=False, gc_chunk_interval=10)[source]

Bases: object

Configuration for memory management in large dataset fitting.

memory_limit_gb

Maximum memory to use in GB (default: 8.0)

Type:

float

safety_factor

Safety factor for memory calculations (default: 0.8)

Type:

float

min_chunk_size

Minimum chunk size in data points (default: 1000)

Type:

int

max_chunk_size

Maximum chunk size in data points (default: 1000000)

Type:

int

use_streaming

Use adaptive hybrid streaming optimization for unlimited data (default: True)

Type:

bool

streaming_batch_size

Chunk size for adaptive hybrid streaming (default: 50000)

Type:

int

streaming_max_epochs

Maximum Gauss-Newton iterations for adaptive hybrid streaming (default: 10)

Type:

int

min_success_rate

Minimum success rate for chunked fitting (default: 0.5) If success rate falls below this threshold, fitting is considered failed

Type:

float

save_diagnostics

Whether to compute and save detailed diagnostic statistics (default: False) When False, skips statistical computations for successful chunks (5-10% faster)

Type:

bool

gc_chunk_interval

Chunks between gc.collect() calls (default: 10, FR-007) Controls how often garbage collection runs during chunked processing. Higher values reduce GC overhead but may increase memory usage.

Type:

int

memory_limit_gb: float
safety_factor: float
min_chunk_size: int
max_chunk_size: int
use_streaming: bool
streaming_batch_size: int
streaming_max_epochs: int
min_success_rate: float
save_diagnostics: bool
gc_chunk_interval: int
__init__(memory_limit_gb=8.0, safety_factor=0.8, min_chunk_size=1000, max_chunk_size=1000000, use_streaming=True, streaming_batch_size=50000, streaming_max_epochs=10, min_success_rate=0.5, save_diagnostics=False, gc_chunk_interval=10)
class nlsq.streaming.large_dataset.LargeDatasetFitter(memory_limit_gb=8.0, config=None, curve_fit_class=None, logger=None, multistart=False, n_starts=10, sampler='lhs')[source]

Bases: object

Large dataset curve fitting with automatic memory management and chunking.

This class handles datasets with millions to billions of points that exceed available memory through automatic chunking, progressive parameter refinement, and streaming optimization. It maintains fitting accuracy while preventing memory overflow through dynamic memory monitoring and chunk size optimization.

Core Capabilities

  • Automatic memory estimation based on data size and parameter count

  • Dynamic chunk size calculation considering available system memory

  • Sequential parameter refinement across data chunks with convergence tracking

  • Streaming optimization for unlimited datasets (no accuracy loss)

  • Real-time progress monitoring with ETA for long-running fits

  • Full integration with NLSQ optimization algorithms and GPU acceleration

  • Multi-start optimization for global search (uses full data)

Memory Management Algorithm

  1. Estimates total memory requirements from dataset size and parameter count

  2. Calculates optimal chunk sizes considering available memory and safety margins

  3. Monitors actual memory usage during processing to prevent overflow

  4. Uses streaming optimization for extremely large datasets (processes all data)

Processing Strategies

  • Single Pass: For datasets fitting within memory limits

  • Sequential Chunking: Processes data in optimal-sized chunks with parameter propagation

  • Streaming Optimization: Mini-batch gradient descent for unlimited datasets (no subsampling)

Multi-Start Optimization

For medium-sized datasets (1M-100M points), multi-start optimization explores multiple starting points on full data, and the best starting point is then used for the full chunked optimization.

Performance Characteristics

  • Maintains <1% parameter error for well-conditioned problems using chunking

  • Achieves 5-50x speedup over naive approaches through memory optimization

  • Scales to datasets of unlimited size using streaming (processes all data)

  • Provides linear time complexity with respect to chunk count

Model Validation Caching (Task Group 7 - 5.1a)

Model functions are validated once per unique function identity using a cache keyed by (id(func), id(func.__code__)). This avoids redundant validation across chunks, providing 1-5% performance gain in chunked processing.

param memory_limit_gb:

Maximum memory usage in GB. System memory is auto-detected if None.

type memory_limit_gb:

float, default 8.0

param config:

Advanced configuration for fine-tuning memory management behavior.

type config:

LDMemoryConfig, optional

param curve_fit_class:

Custom CurveFit instance for specialized fitting requirements.

type curve_fit_class:

nlsq.minpack.CurveFit, optional

param multistart:

Enable multi-start optimization for global search.

type multistart:

bool, default False

param n_starts:

Number of starting points for multi-start optimization.

type n_starts:

int, default 10

param sampler:

Sampling strategy for multi-start: ‘lhs’, ‘sobol’, or ‘halton’.

type sampler:

str, default ‘lhs’

config

Active memory management configuration

Type:

LDMemoryConfig

curve_fitter

Internal curve fitting engine with JAX acceleration

Type:

nlsq.minpack.CurveFit

logger

Internal logging for performance monitoring and debugging

Type:

Logger

fit : Main fitting method with automatic memory management
fit_with_progress : Fitting with real-time progress reporting and ETA
get_memory_recommendations : Pre-fitting memory analysis and strategy recommendations
Important: Chunking-Compatible Model Functions
-----------------------------------------------
When using chunked processing (for datasets > memory limit), your model function
MUST respect the size of xdata. During chunking, xdata will be a subset of the
full dataset, and your model must return output matching that subset size.
\*\*INCORRECT - Model ignores xdata size (will cause shape mismatch errors):**
>>> def bad_model(xdata, a, b):
...     # WRONG: Always returns full array, ignoring xdata size
...     t_full = jnp.arange(10_000_000)  # Fixed size!
...     return a * jnp.exp(-b * t_full)  # Shape mismatch during chunking
\*\*CORRECT - Model respects xdata size:**
>>> def good_model(xdata, a, b):
...     # CORRECT: Uses xdata as indices to return only requested subset
...     indices = xdata.astype(jnp.int32)
...     return a * jnp.exp(-b * indices)  # Shape matches xdata
\*\*Alternative - Direct computation on xdata:**
>>> def direct_model(xdata, a, b):
...     # CORRECT: Operates directly on xdata
...     return a * jnp.exp(-b * xdata)  # Shape automatically matches

Examples

Basic usage with automatic configuration:

>>> import numpy as np
>>> import jax.numpy as jnp
>>>
>>> # 10 million data points
>>> x = np.linspace(0, 10, 10_000_000)
>>> y = 2.5 * jnp.exp(-1.3 * x) + 0.1 + np.random.normal(0, 0.05, len(x))
>>>
>>> fitter = LargeDatasetFitter(memory_limit_gb=4.0)
>>> result = fitter.fit(
...     lambda x, a, b, c: a * jnp.exp(-b * x) + c,
...     x, y, p0=[2, 1, 0]
... )
>>> print(f"Parameters: {result.popt}")
>>> print(f"Chunks used: {result.n_chunks}")

Multi-start optimization:

>>> fitter = LargeDatasetFitter(
...     memory_limit_gb=4.0,
...     multistart=True,
...     n_starts=10,
...     sampler='lhs',
... )
>>> result = fitter.fit(
...     lambda x, a, b, c: a * jnp.exp(-b * x) + c,
...     x, y, p0=[2, 1, 0],
...     bounds=([0, 0, 0], [10, 5, 10])
... )

Advanced configuration with progress monitoring:

>>> config = LDMemoryConfig(
...     memory_limit_gb=8.0,
...     min_chunk_size=10000,
...     max_chunk_size=1000000,
...     use_streaming=True,
...     streaming_batch_size=50000
... )
>>> fitter = LargeDatasetFitter(config=config)
>>>
>>> # Fit with progress bar for long-running operation
>>> result = fitter.fit_with_progress(
...     exponential_model, x_huge, y_huge, p0=[2, 1, 0]
... )

Memory analysis before processing:

>>> recommendations = fitter.get_memory_recommendations(len(x), n_params=3)
>>> print(f"Strategy: {recommendations['processing_strategy']}")
>>> print(f"Memory estimate: {recommendations['memory_estimate_gb']:.2f} GB")
>>> print(f"Recommended chunks: {recommendations['n_chunks']}")

See also

curve_fit_large

High-level function with automatic dataset size detection

LDMemoryConfig

Configuration class for memory management parameters

estimate_memory_requirements

Standalone function for memory estimation

Notes

The sequential chunking algorithm maintains parameter accuracy by using each chunk’s result as the initial guess for the next chunk. This approach typically maintains fitting accuracy within 0.1% of single-pass results for well-conditioned problems while enabling processing of arbitrarily large datasets.

For extremely large datasets, streaming optimization processes all data using mini-batch gradient descent with no subsampling, ensuring zero accuracy loss compared to subsampling approaches (removed in v0.2.0).

__init__(memory_limit_gb=8.0, config=None, curve_fit_class=None, logger=None, multistart=False, n_starts=10, sampler='lhs')[source]

Initialize LargeDatasetFitter.

Parameters:

  • memory_limit_gb (float, optional) – Memory limit in GB (default: 8.0)

  • config (LDMemoryConfig, optional) – Custom memory configuration

  • curve_fit_class (nlsq.minpack.CurveFit, optional) – Custom CurveFit instance to use

  • logger (logging.Logger, optional) – External logger instance for integration with application logging. If None, uses NLSQ’s internal logger. This allows chunk failure warnings to appear in your application’s logs.

  • multistart (bool, optional) – Enable multi-start optimization for global search (default: False). When enabled, explores multiple starting points on full data before running the full chunked optimization.

  • n_starts (int, optional) – Number of starting points for multi-start optimization (default: 10). Set to 0 to disable multi-start even when multistart=True.

  • sampler (str, optional) – Sampling strategy for generating starting points (default: ‘lhs’). Options: ‘lhs’ (Latin Hypercube), ‘sobol’, ‘halton’.

last_stats: DatasetStats | None
fit_history: list[dict]
estimate_requirements(n_points, n_params)[source]

Estimate memory requirements and processing strategy.

Parameters:

  • n_points (int) – Number of data points

  • n_params (int) – Number of parameters to fit

Returns:

Detailed statistics and recommendations

Return type:

DatasetStats

fit(f, xdata, ydata, p0=None, bounds=(-inf, inf), method='trf', solver='auto', **kwargs)[source]

Fit curve to large dataset with automatic memory management.

Parameters:

  • f (callable) – The model function f(x, *params) -> y

  • xdata (np.ndarray) – Independent variable data

  • ydata (np.ndarray) – Dependent variable data

  • p0 (array-like, optional) – Initial parameter guess

  • bounds (tuple, optional) – Parameter bounds (lower, upper)

  • method (str, optional) – Optimization method (default: ‘trf’)

  • solver (str, optional) – Solver type (default: ‘auto’)

  • **kwargs – Additional arguments passed to curve_fit

Returns:

Optimization result with fitted parameters and statistics

Return type:

OptimizeResult

fit_with_progress(f, xdata, ydata, p0=None, bounds=(-inf, inf), method='trf', solver='auto', **kwargs)[source]

Fit curve with progress reporting for long-running fits.

Parameters:

  • f (callable) – The model function f(x, *params) -> y

  • xdata (np.ndarray) – Independent variable data

  • ydata (np.ndarray) – Dependent variable data

  • p0 (array-like, optional) – Initial parameter guess

  • bounds (tuple, optional) – Parameter bounds (lower, upper)

  • method (str, optional) – Optimization method (default: ‘trf’)

  • solver (str, optional) – Solver type (default: ‘auto’)

  • **kwargs – Additional arguments passed to curve_fit

Returns:

Optimization result with fitted parameters and statistics

Return type:

OptimizeResult

memory_monitor()[source]

Context manager for monitoring memory usage during fits.

get_memory_recommendations(n_points, n_params)[source]

Get memory usage recommendations for a dataset.

Parameters:

  • n_points (int) – Number of data points

  • n_params (int) – Number of parameters

Returns:

Recommendations and memory analysis

Return type:

dict

class nlsq.streaming.large_dataset.MemoryEstimator[source]

Bases: object

Utilities for estimating memory usage and optimal chunk sizes.

This class provides CPU memory detection via psutil, with fallback to 16GB when detection fails (e.g., containerized environments with cgroups).

static estimate_memory_per_point(n_params, use_jacobian=True)[source]

Estimate memory usage per data point in bytes.

Parameters:

  • n_params (int) – Number of parameters

  • use_jacobian (bool, optional) – Whether Jacobian computation is needed (default: True)

Returns:

Estimated memory usage per point in bytes

Return type:

float

static get_available_memory_gb()[source]

Get available system memory in GB.

Returns:

Available memory in GB. Falls back to 16GB if detection fails.

Return type:

float

Notes

  • Re-evaluates on each call (no caching) per requirements.

  • Falls back to 16GB when detection fails (containerized environments).

static get_total_available_memory_gb()[source]

Get total available memory (CPU + GPU) in GB.

Combines available CPU memory from psutil with available GPU memory from JAX device API. Re-evaluates on each call (no caching).

Returns:

Total available memory (CPU + GPU) in GB.

Return type:

float

Notes

  • CPU memory: Uses psutil.virtual_memory().available

  • GPU memory: Uses JAX device API via GPUMemoryEstimator

  • Falls back to 16GB for CPU if detection fails

  • Returns 0 for GPU if detection fails or no GPU present

static estimate_maximum_memory_usage_gb(n_points, n_params, safety_factor=1.2)[source]

Estimate maximum memory usage to prevent crashes.

Parameters:

  • n_points (int) – Number of data points

  • n_params (int) – Number of parameters

  • safety_factor (float, optional) – Safety factor for memory estimation (default: 1.2)

Returns:

Estimated maximum memory usage in GB.

Return type:

float

Notes

This method estimates the peak memory usage during optimization, which is useful for workflow selection to prevent out-of-memory crashes.

static calculate_optimal_chunk_size(n_points, n_params, memory_config)[source]

Calculate optimal chunk size based on memory constraints.

Parameters:

  • n_points (int) – Total number of data points

  • n_params (int) – Number of parameters

  • memory_config (LDMemoryConfig) – Memory configuration

Returns:

Optimal chunk size and dataset statistics

Return type:

tuple[int, DatasetStats]

class nlsq.streaming.large_dataset.ProgressReporter(total_chunks, logger=None)[source]

Bases: object

Progress reporting for long-running fits.

__init__(total_chunks, logger=None)[source]

Initialize progress reporter.

Parameters:

  • total_chunks (int) – Total number of chunks to process

  • logger (optional) – Logger instance for reporting progress

update(chunk_idx, chunk_result=None)[source]

Update progress.

Parameters:

  • chunk_idx (int) – Index of completed chunk

  • chunk_result (dict, optional) – Results from chunk processing

nlsq.streaming.large_dataset.cleanup_memory()[source]

Perform memory cleanup between workflow phases.

This function clears both Python garbage and JAX compilation caches, designed to be called between workflow phases to free memory.

Notes

  • Calls gc.collect() to trigger Python garbage collection

  • Calls jax.clear_caches() to clear JAX JIT compilation caches

  • Handles errors gracefully (does not raise exceptions)

Examples

>>> from nlsq.streaming.large_dataset import cleanup_memory
>>> # After completing a workflow phase
>>> cleanup_memory()
nlsq.streaming.large_dataset.estimate_memory_requirements(n_points, n_params)[source]

Estimate memory requirements for a dataset.

Parameters:

  • n_points (int) – Number of data points

  • n_params (int) – Number of parameters

Returns:

Memory requirements and processing recommendations

Return type:

DatasetStats

Examples

>>> from nlsq.streaming.large_dataset import estimate_memory_requirements
>>>
>>> # Estimate requirements for 50M points, 3 parameters
>>> stats = estimate_memory_requirements(50_000_000, 3)
>>> print(f"Estimated memory: {stats.total_memory_estimate_gb:.2f} GB")
>>> print(f"Recommended chunk size: {stats.recommended_chunk_size:,}")
>>> print(f"Number of chunks: {stats.n_chunks}")
nlsq.streaming.large_dataset.fit_large_dataset(f, xdata, ydata, p0=None, memory_limit_gb=8.0, show_progress=False, logger=None, multistart=False, n_starts=10, sampler='lhs', **kwargs)[source]

Convenience function for fitting large datasets.

Parameters:

  • f (callable) – The model function f(x, *params) -> y

  • xdata (np.ndarray) – Independent variable data

  • ydata (np.ndarray) – Dependent variable data

  • p0 (array-like, optional) – Initial parameter guess

  • memory_limit_gb (float, optional) – Memory limit in GB (default: 8.0)

  • show_progress (bool, optional) – Whether to show progress (default: False)

  • logger (logging.Logger, optional) – External logger for application integration (default: None)

  • multistart (bool, optional) – Enable multi-start optimization for global search (default: False). When enabled, explores multiple starting points on full data before running the full chunked optimization.

  • n_starts (int, optional) – Number of starting points for multi-start optimization (default: 10). Set to 0 to disable multi-start even when multistart=True.

  • sampler (str, optional) – Sampling strategy for generating starting points (default: ‘lhs’). Options: ‘lhs’ (Latin Hypercube), ‘sobol’, ‘halton’.

  • **kwargs – Additional arguments passed to curve_fit

Returns:

Optimization result

Return type:

OptimizeResult

Examples

>>> from nlsq.streaming.large_dataset import fit_large_dataset
>>> import numpy as np
>>> import jax.numpy as jnp
>>>
>>> # Generate large dataset
>>> x_large = np.linspace(0, 10, 5_000_000)
>>> y_large = 2.5 * np.exp(-1.3 * x_large) + np.random.normal(0, 0.1, len(x_large))
>>>
>>> # Fit with automatic memory management
>>> result = fit_large_dataset(
...     lambda x, a, b: a * jnp.exp(-b * x),
...     x_large, y_large,
...     p0=[2.0, 1.0],
...     memory_limit_gb=4.0,
...     show_progress=True
... )
>>> print(f"Fitted parameters: {result.popt}")
>>> print(f"Success rate: {result.success_rate:.1%}")
>>>
>>> # Fit with multi-start optimization
>>> result = fit_large_dataset(
...     lambda x, a, b: a * jnp.exp(-b * x),
...     x_large, y_large,
...     p0=[2.0, 1.0],
...     bounds=([0, 0], [10, 5]),
...     multistart=True,
...     n_starts=10,
...     sampler='lhs'
... )
>>>
>>> # Check failure diagnostics if some chunks failed
>>> if result.failure_summary['total_failures'] > 0:
...     print(f"Failed chunks: {result.failure_summary['failed_chunk_indices']}")
...     print(f"Common errors: {result.failure_summary['common_errors']}")
nlsq.streaming.large_dataset.get_cached_devices()[source]

Get cached JAX devices, computing once on first call.

Overview

The nlsq.large_dataset module provides specialized tools for fitting curves to datasets that are too large to fit in memory or require chunking for efficient processing. This module is essential for working with datasets containing millions or billions of points.

Key Features

  • Automatic dataset handling for 100M+ points

  • Intelligent chunking with <1% error for well-conditioned problems

  • Memory estimation and automatic memory management

  • Streaming optimization for unlimited-size datasets

  • Progress reporting for long-running fits

Classes

LargeDatasetFitter([memory_limit_gb, ...])

Large dataset curve fitting with automatic memory management and chunking.
class nlsq.streaming.large_dataset.LDMemoryConfig(memory_limit_gb=8.0, safety_factor=0.8, min_chunk_size=1000, max_chunk_size=1000000, use_streaming=True, streaming_batch_size=50000, streaming_max_epochs=10, min_success_rate=0.5, save_diagnostics=False, gc_chunk_interval=10)[source]

Bases: object

Configuration for memory management in large dataset fitting.

memory_limit_gb

Maximum memory to use in GB (default: 8.0)

Type:

float

safety_factor

Safety factor for memory calculations (default: 0.8)

Type:

float

min_chunk_size

Minimum chunk size in data points (default: 1000)

Type:

int

max_chunk_size

Maximum chunk size in data points (default: 1000000)

Type:

int

use_streaming

Use adaptive hybrid streaming optimization for unlimited data (default: True)

Type:

bool

streaming_batch_size

Chunk size for adaptive hybrid streaming (default: 50000)

Type:

int

streaming_max_epochs

Maximum Gauss-Newton iterations for adaptive hybrid streaming (default: 10)

Type:

int

min_success_rate

Minimum success rate for chunked fitting (default: 0.5) If success rate falls below this threshold, fitting is considered failed

Type:

float

save_diagnostics

Whether to compute and save detailed diagnostic statistics (default: False) When False, skips statistical computations for successful chunks (5-10% faster)

Type:

bool

gc_chunk_interval

Chunks between gc.collect() calls (default: 10, FR-007) Controls how often garbage collection runs during chunked processing. Higher values reduce GC overhead but may increase memory usage.

Type:

int

memory_limit_gb: float
safety_factor: float
min_chunk_size: int
max_chunk_size: int
use_streaming: bool
streaming_batch_size: int
streaming_max_epochs: int
min_success_rate: float
save_diagnostics: bool
gc_chunk_interval: int
__init__(memory_limit_gb=8.0, safety_factor=0.8, min_chunk_size=1000, max_chunk_size=1000000, use_streaming=True, streaming_batch_size=50000, streaming_max_epochs=10, min_success_rate=0.5, save_diagnostics=False, gc_chunk_interval=10)

Functions

fit_large_dataset(f, xdata, ydata[, p0, ...])

Convenience function for fitting large datasets.

estimate_memory_requirements(n_points, n_params)

Estimate memory requirements for a dataset.

Examples

Basic Usage with curve_fit_large

from nlsq import curve_fit_large, estimate_memory_requirements
import jax.numpy as jnp
import numpy as np

# Check memory requirements
n_points = 50_000_000  # 50 million points
n_params = 3
stats = estimate_memory_requirements(n_points, n_params)
print(f"Memory required: {stats.total_memory_estimate_gb:.2f} GB")

# Generate large dataset
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)


# Define fit function
def exponential(x, a, b, c):
    return a * jnp.exp(-b * x) + c


# Fit with automatic chunking
popt, pcov = curve_fit_large(
    exponential,
    x,
    y,
    p0=[2.5, 0.6, 0.2],
    memory_limit_gb=4.0,
    show_progress=True,
)

Advanced Usage with LargeDatasetFitter

from nlsq import LargeDatasetFitter, LDMemoryConfig
import jax.numpy as jnp

# Configure memory management
config = LDMemoryConfig(
    memory_limit_gb=4.0,
    min_chunk_size=10000,
    max_chunk_size=1000000,
    min_success_rate=0.8,  # Require 80% of chunks to succeed
)

# Create fitter
fitter = LargeDatasetFitter(config=config)

# Fit with progress tracking
result = fitter.fit_with_progress(exponential, x, y, p0=[2.5, 0.6, 0.2])

print(f"Fitted parameters: {result.popt}")
# Note: success_rate and n_chunks only available for multi-chunk fits

Adaptive Hybrid Streaming

For datasets that don’t fit in memory, use adaptive hybrid streaming:

from nlsq import AdaptiveHybridStreamingOptimizer, HybridStreamingConfig

config = HybridStreamingConfig(chunk_size=10000, gauss_newton_max_iterations=100)
optimizer = AdaptiveHybridStreamingOptimizer(config)

result = optimizer.fit((x, y), func, p0=p0)

Sparse Jacobian Optimization

For problems with sparse structure:

from nlsq import SparseJacobianComputer

# Detect and exploit sparsity
sparse_computer = SparseJacobianComputer(sparsity_threshold=0.01)
pattern, sparsity = sparse_computer.detect_sparsity_pattern(func, p0, x_sample)

if sparsity > 0.1:  # If more than 10% sparse
    print(f"Jacobian is {sparsity:.1%} sparse")

See Also