nlsq.minpack module

This module provides the main user-facing API with the CurveFit class and curve_fit function.

MINPACK-style algorithms for nonlinear least squares optimization.

This module provides JAX implementations of classic MINPACK algorithms including Levenberg-Marquardt and Trust Region Reflective methods with GPU/TPU acceleration. These algorithms are the foundation of NLSQ’s curve fitting capabilities.

The MINPACK algorithms combine robust optimization strategies with efficient linear algebra operations, all JIT-compiled for high performance on modern hardware.

Key Components:

  • fit: Unified entry point with workflow-based automatic strategy selection

  • curve_fit: Main high-level interface (SciPy-compatible)

  • CurveFit: Class-based interface with state management

  • Levenberg-Marquardt: Classic damped least squares algorithm

  • Trust Region: Advanced bounded optimization

Algorithms:

  • ‘lm’: Levenberg-Marquardt (unbounded problems)

  • ‘trf’: Trust Region Reflective (bounded problems, recommended)

  • ‘dogbox’: Dogleg with rectangular trust regions

Example

>>> from nlsq import CurveFit
>>> import jax.numpy as jnp
>>>
>>> def exponential(x, a, b): return a * jnp.exp(-b * x)
>>>
>>> # Class-based interface for reusing compilations
>>> fitter = CurveFit()
>>> popt1, pcov1 = fitter.curve_fit(exponential, x1, y1, p0=[2.0, 0.5])
>>> popt2, pcov2 = fitter.curve_fit(exponential, x2, y2, p0=[2.5, 0.6])

See also

nlsq.curve_fit : Function-based interface nlsq.fit : Unified entry point with automatic workflow selection nlsq.least_squares : Lower-level optimization interface nlsq.trf : Trust Region Reflective implementation

class nlsq.core.minpack.CurveFit(flength=None, use_dynamic_sizing=False, enable_stability=False, enable_recovery=False, enable_overflow_check=False, cache_config=None, max_jacobian_elements_for_svd=10000000)[source]

Bases: object

Main class for nonlinear least squares curve fitting with JAX acceleration.

This class provides the core curve fitting functionality with JAX JIT compilation, automatic differentiation for Jacobian computation, and multiple optimization algorithms. It handles data preprocessing, optimization algorithm selection, and covariance matrix computation.

The class maintains compiled versions of fitting functions to avoid recompilation overhead when fitting multiple datasets with the same function signature.

flength

Fixed data length for input padding to avoid JAX retracing.

Type:

float or None

use_dynamic_sizing

Whether to use dynamic sizing instead of fixed padding.

Type:

bool

logger

Internal logger for debugging and performance monitoring.

Type:

Logger

curve_fit : Main fitting method
create_sigma_transform_funcs : Internal method for sigma transformation setup
__init__(flength=None, use_dynamic_sizing=False, enable_stability=False, enable_recovery=False, enable_overflow_check=False, cache_config=None, max_jacobian_elements_for_svd=10000000)[source]

Initialize CurveFit instance.

Parameters:

  • flength (float, optional) – Fixed data length for JAX compilation. Input data is padded to this length to avoid recompilation when fitting datasets of different sizes. If None, no padding is applied and each dataset size triggers recompilation. Ignored when use_dynamic_sizing=True for large datasets.

  • use_dynamic_sizing (bool, default False) – Enable dynamic sizing to reduce memory usage. When True, padding is only applied when data size is smaller than flength. For large datasets, uses actual size to prevent excessive memory allocation. Default False maintains backward compatibility with fixed padding behavior.

  • enable_stability (bool, default False) – Enable numerical stability checks and fixes (validation, algorithm selection). Note: This does NOT include overflow checking which adds overhead.

  • enable_recovery (bool, default False) – Enable automatic recovery from optimization failures.

  • enable_overflow_check (bool, default False) – Enable overflow/underflow checking in function evaluations. This adds ~30% overhead so it’s separate from other stability features.

  • cache_config (dict, optional) – Configuration for the unified JIT compilation cache. Supported keys: 'maxsize' (int, default 128) — maximum number of compiled functions to cache; 'enable_stats' (bool, default True) — track cache statistics (hits, misses, compile_time_ms); 'disk_cache_enabled' (bool, default False) — enable disk caching tier (Phase 2 feature). If None, uses global cache with default settings.

  • max_jacobian_elements_for_svd (int, default 10_000_000) – Maximum Jacobian size (m x n elements) for SVD computation during stability checks. SVD is skipped for larger Jacobians to avoid O(min(m,n)^2 x max(m,n)) overhead. Only applies when enable_stability=True.

Notes

Fixed length compilation trades memory usage for compilation speed: - flength=None: Minimal memory, recompiles for each dataset size - flength=large_value: Higher memory, avoids recompilation - use_dynamic_sizing=True: Balanced approach for mixed dataset sizes

update_flength(flength)[source]

Set the fixed input data length.

Parameters:

flength (float) – The fixed input data length.

create_sigma_transform_funcs()[source]

Create JIT-compiled sigma transform functions.

This function creates two JIT-compiled functions: sigma_transform1d and sigma_transform2d, which are used to compute the sigma transform for 1D and 2D data, respectively. The functions are stored as attributes of the object on which the method is called.

create_covariance_svd()[source]

Create JIT-compiled SVD function for covariance computation.

pad_fit_data(xdata, ydata, xdims, len_diff)[source]

Pad fit data to match the fixed input data length.

This function pads the input data arrays with small values to match the fixed input data length to avoid JAX retracing the JITted functions. The padding is added along the second dimension of the xdata array if it’s multidimensional data otherwise along the first dimension. The small values are chosen to be EPS, a global constant defined as a very small positive value which avoids numerical issues.

Parameters:

  • xdata (np.ndarray) – The independent variables of the data.

  • ydata (np.ndarray) – The dependent variables of the data.

  • xdims (int) – The number of dimensions in the xdata array.

  • len_diff (int) – The difference in length between the data arrays and the fixed input data length.

Returns:

The padded xdata and ydata arrays.

Return type:

Tuple[np.ndarray, np.ndarray]

curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False, check_finite=True, bounds=(-inf, inf), method=None, solver='auto', batch_size=None, jac=None, data_mask=None, timeit=False, return_eval=False, full_output=False, callback=None, compute_diagnostics=False, diagnostics_level=DiagnosticLevel.BASIC, diagnostics_config=None, **kwargs)[source]

Use non-linear least squares to fit a function, f, to data. Assumes ydata = f(xdata, \*params) + eps.

Parameters:

method (str | None, optional) – Optimization algorithm. Options: - ‘trf’ (default): Trust Region Reflective - ‘lm’: Levenberg-Marquardt - ‘hybrid_streaming’: Adaptive Hybrid Streaming Optimizer (for large datasets) - ‘auto’: Memory-based automatic selection (streaming/chunked/standard) If None, auto-selects ‘trf’.

nlsq.core.minpack.curve_fit(f, xdata, ydata, *args, auto_bounds=False, bounds_safety_factor=10.0, stability=False, rescale_data=True, max_jacobian_elements_for_svd=10000000, fallback=False, max_fallback_attempts=10, fallback_verbose=False, multistart=False, n_starts=10, global_search=False, sampler='lhs', center_on_p0=True, scale_factor=1.0, method=None, cmaes_config=None, compute_diagnostics=False, diagnostics_level=DiagnosticLevel.BASIC, diagnostics_config=None, **kwargs)[source]

Use nonlinear least squares to fit a function to data with GPU/TPU acceleration.

This is the main user-facing function that provides a drop-in replacement for scipy.optimize.curve_fit with GPU/TPU acceleration via JAX. The function automatically handles JAX JIT compilation, double precision configuration, and optimization algorithm selection.

Parameters:

  • f (callable) – The model function f(x, *popt) -> y. Must be JAX-compatible, meaning it should use jax.numpy instead of numpy for mathematical operations to enable GPU acceleration and automatic differentiation.

  • xdata (array_like) – The independent variable where the data is measured.

  • ydata (array_like) – The dependent data, nominally f(xdata, *popt).

  • auto_bounds (bool, optional) –

    Enable automatic parameter bounds inference from data characteristics. When True, reasonable bounds are inferred based on:

    • Data ranges (x and y)

    • Initial parameter guess (p0)

    • Parameter positivity constraints

    • Safety factors to avoid over-constraining

    The inferred bounds are merged with any user-provided bounds via the bounds parameter. User bounds take precedence where specified. Default: False.

  • bounds_safety_factor (float, optional) – Safety multiplier for automatic bounds (larger = more conservative). Only used when auto_bounds=True. Default: 10.0.

  • stability ({'auto', 'check', False}, optional) –

    Control numerical stability checks and automatic fixes:

    • ’auto’: Check for stability issues and automatically apply fixes (optionally rescale data, normalize parameters, handle NaN/Inf)

    • ’check’: Check for stability issues and warn, but don’t apply fixes

    • False: Skip stability checks entirely (default)

    When ‘auto’, detected issues are fixed before optimization:

    • Ill-conditioned data (condition number > 1e10) is rescaled to [0, 1] (only if rescale_data=True)

    • Large data ranges (> 1e4) are normalized (only if rescale_data=True)

    • NaN/Inf values are replaced with mean

    • Parameter scale mismatches (ratio > 1e6) are normalized

    Default: False.

  • rescale_data (bool, optional) – When stability=’auto’, controls whether data is automatically rescaled to [0, 1] for ill-conditioned or large-range data. Set to False for applications where data must maintain physical units (e.g., time in seconds, frequency in Hz). NaN/Inf handling and parameter normalization are still applied when stability=’auto’. Default: True.

  • max_jacobian_elements_for_svd (int, optional) –

    Maximum number of elements in the Jacobian matrix (m x n) for which SVD will be computed during stability checks. For larger Jacobians, SVD is skipped to avoid O(min(m,n)^2 x max(m,n)) computation overhead. NaN/Inf checking is still performed for large Jacobians.

    Examples of element counts: - 1M data points x 7 params = 7M elements - 100K data points x 100 params = 10M elements

    Set to a larger value if you need condition number checks for large problems, or a smaller value to skip SVD more aggressively. Default: 10,000,000 (10M elements).

  • fallback (bool, optional) –

    Enable automatic fallback strategies for difficult optimization problems. When True, the optimizer will automatically try alternative approaches if the initial optimization fails, including:

    • Alternative optimization methods

    • Perturbed initial guesses

    • Relaxed tolerances

    • Inferred parameter bounds

    • Robust loss functions

    • Problem rescaling

    Default: False. Enabling this improves success rate on difficult problems but adds overhead when optimizations fail.

  • max_fallback_attempts (int, optional) – Maximum number of fallback attempts to try before giving up. Only used when fallback=True. Default: 10.

  • fallback_verbose (bool, optional) – Print detailed information about fallback attempts. Only used when fallback=True. Default: False.

  • multistart (bool, optional) – Enable multi-start optimization for global search. When True, generates multiple starting points using Latin Hypercube Sampling (or other samplers) and evaluates each, selecting the best result. This helps find global optima in problems with multiple local minima. Default: False.

  • n_starts (int, optional) – Number of starting points for multi-start optimization. Only used when multistart=True or global_search=True. Default: 10.

  • global_search (bool, optional) – Shorthand for enabling multi-start with n_starts=20. Equivalent to multistart=True, n_starts=20. Useful for thorough global search. Default: False.

  • sampler ({'lhs', 'sobol', 'halton'}, optional) –

    Sampling strategy for generating starting points in multi-start:

    • ’lhs’: Latin Hypercube Sampling (stratified random, default)

    • ’sobol’: Sobol quasi-random sequence (deterministic, low-discrepancy)

    • ’halton’: Halton quasi-random sequence (deterministic, prime bases)

    Only used when multistart=True or global_search=True. Default: ‘lhs’.

  • center_on_p0 (bool, optional) – When True, center multi-start samples around the initial guess p0 rather than uniformly across the full parameter bounds. This provides more focused exploration around a data-informed starting region. Only used when multistart=True or global_search=True. Default: True.

  • scale_factor (float, optional) – Scale factor for the exploration region when center_on_p0=True. Multiplier for the exploration range around p0. Smaller values (0.5) mean tighter exploration, larger values (2.0) mean wider exploration. Only used when multistart=True or global_search=True. Default: 1.0.

  • method ({'auto', 'cmaes', 'multi-start', 'trf'} | None, optional) –

    Optimization method to use:

    • ’auto’: Automatically select based on parameter scale ratio. Uses CMA-ES for multi-scale problems (>1000x scale difference) when evosax is installed, otherwise multi-start.

    • ’cmaes’: Use CMA-ES (Covariance Matrix Adaptation Evolution Strategy) for gradient-free global optimization. Requires bounds. Best for multi-scale parameters spanning many orders of magnitude. Falls back to multi-start if evosax is not installed.

    • ’multi-start’: Use multi-start optimization with Latin Hypercube Sampling for initial points.

    • ’trf’: Use Trust Region Reflective (default behavior).

    • None: Default behavior (TRF with optional multi-start if enabled).

    Default: None.

  • cmaes_config (CMAESConfig | None, optional) – Configuration for CMA-ES optimization. If None, uses default config or preset specified by method parameter. See CMAESConfig for options including max_generations, restart_strategy, and popsize. Only used when method=’cmaes’ or method=’auto’ selects CMA-ES. Default: None.

  • *args (Any) – Additional arguments passed to CurveFit.curve_fit method.

  • **kwargs (Any) – Additional arguments passed to CurveFit.curve_fit method.

Returns:

  • popt (ndarray) – Optimal values for the parameters.

  • pcov (ndarray) – The estimated covariance of popt.

  • When fallback=True, the returned object also contains

  • - fallback_strategy_used (str or None) – Name of the fallback strategy that succeeded, or None if original succeeded

  • - fallback_attempts (int) – Number of optimization attempts before success

  • When multistart=True or global_search=True, the returned object contains

  • - multistart_diagnostics (dict) – Dictionary with multi-start exploration details including n_starts, best_loss, all_losses, exploration time, etc.

Return type:

tuple[np.ndarray, np.ndarray] | CurveFitResult

Notes

This function creates a CurveFit instance internally and calls its curve_fit method. For multiple fits with the same function signature, consider creating a CurveFit instance directly to benefit from JAX compilation caching.

When fallback=True, the optimizer tries increasingly aggressive recovery strategies if the initial optimization fails. This is particularly useful for:

  • Poor initial parameter guesses

  • Ill-conditioned problems

  • Problems with outliers

  • Numerically challenging models

When multistart=True or global_search=True, the optimizer explores multiple starting points to find the global optimum. This is particularly useful for:

  • Problems with multiple local minima

  • Complex multi-modal objective functions

  • Cases where the initial guess may not be close to the global optimum

See also

CurveFit.curve_fit

The underlying method with full parameter documentation

fit

Unified entry point with automatic workflow selection

curve_fit_large

For datasets with millions of points requiring special handling

FallbackOrchestrator

Direct access to fallback system for custom configurations

MultiStartOrchestrator

Direct access to multi-start system

Examples

Basic usage without fallback:

>>> import jax.numpy as jnp
>>> import numpy as np
>>>
>>> def exponential(x, a, b):
...     return a * jnp.exp(-b * x)
>>>
>>> x = np.linspace(0, 4, 50)
>>> y = 2.5 * np.exp(-1.3 * x) + 0.1 * np.random.normal(size=len(x))
>>> popt, _pcov = curve_fit(exponential, x, y, p0=[2, 1])

Using multi-start for global search:

>>> # Enable multi-start with default n_starts=10
>>> result = curve_fit(exponential, x, y, p0=[2, 1],
...                   bounds=([0, 0], [10, 5]), multistart=True)
>>>
>>> # Use global_search shorthand for thorough exploration (n_starts=20)
>>> result = curve_fit(exponential, x, y, p0=[2, 1],
...                   bounds=([0, 0], [10, 5]), global_search=True)
>>>
>>> # Customize multi-start with different sampler
>>> result = curve_fit(exponential, x, y, p0=[2, 1],
...                   bounds=([0, 0], [10, 5]),
...                   multistart=True, n_starts=15, sampler='sobol')

Using fallback for difficult problems:

>>> # Very poor initial guess - may fail without fallback
>>> result = curve_fit(exponential, x, y, p0=[100, 50], fallback=True)
>>>
>>> # Check which strategy was used
>>> if result.fallback_strategy_used:
...     print(f"Recovered using: {result.fallback_strategy_used}")

Combined multi-start + fallback for maximum robustness:

>>> result = curve_fit(
...     exponential, x, y, p0=[2, 1],
...     bounds=([0, 0], [10, 5]),
...     global_search=True,
...     fallback=True
... )
nlsq.core.minpack.fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False, check_finite=True, bounds=(-inf, inf), method=None, workflow='auto', goal=None, **kwargs)[source]

Unified curve fitting entry point with automatic workflow selection.

This function provides a simplified API for curve fitting that automatically selects the optimal strategy based on dataset size, available memory, and user-specified goals. It coexists with curve_fit() and curve_fit_large() for users who prefer explicit control.

Parameters:

  • f (callable) – Model function f(x, *params) -> y. Must use jax.numpy operations.

  • xdata (array_like) – Independent variable data.

  • ydata (array_like) – Dependent variable data.

  • p0 (array_like, optional) – Initial parameter guess.

  • sigma (array_like, optional) – Uncertainties in ydata for weighted fitting.

  • absolute_sigma (bool, optional) – Whether sigma represents absolute uncertainties.

  • check_finite (bool, optional) – Check for finite input values.

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

  • method (str, optional) – Optimization algorithm (‘trf’, ‘lm’, or None for auto).

  • workflow (str or config object, optional) –

    Workflow selection:

    • ’auto’: Automatically select based on dataset size and memory (default)

    • ’standard’: Basic curve_fit() with defaults

    • ’quality’: Tightest tolerances + multi-start

    • ’fast’: Loose tolerances for speed

    • ’large_robust’: Chunked + multi-start for large datasets

    • ’streaming’: AdaptiveHybridStreamingOptimizer for huge datasets

    • ’hpc_distributed’: Multi-GPU/node configuration

    • ’global_auto’: Auto-selects CMA-ES vs Multi-Start based on scale ratio

    • ’cmaes’: CMA-ES with BIPOP restarts (100 generations)

    • ’cmaes-global’: Thorough CMA-ES exploration (200 gens, 2x population)

    • Custom config object: LDMemoryConfig, HybridStreamingConfig, or GlobalOptimizationConfig

  • goal (str or OptimizationGoal, optional) –

    Optimization goal that modifies tolerances:

    • ’fast’: Prioritize speed, use looser tolerances

    • ’robust’ or ‘global’: Standard tolerances, enable multi-start

    • ’memory_efficient’: Minimize memory usage

    • ’quality’: Highest precision, tighter tolerances, multi-start

  • **kwargs (Any) – Additional optimization parameters (ftol, xtol, gtol, max_nfev, loss). These are passed through to the underlying curve_fit() function.

Returns:

result – Optimization result. Contains popt, pcov, and additional diagnostics. Supports tuple unpacking: popt, pcov = fit(…)

Return type:

CurveFitResult

Raises:

ValueError – If workflow name is invalid or goal string is not recognized.

Examples

Basic usage with automatic workflow selection:

>>> import jax.numpy as jnp
>>> def model(x, a, b): return a * jnp.exp(-b * x)
>>> popt, pcov = fit(model, xdata, ydata, p0=[1, 2])

Using a named workflow:

>>> result = fit(model, xdata, ydata, p0=[1, 2], workflow='quality',
...              bounds=([0, 0], [10, 10]))

Using a goal to adjust tolerances:

>>> result = fit(model, xdata, ydata, p0=[1, 2], goal='quality')

Using a custom config object:

>>> from nlsq.streaming.large_dataset import LDMemoryConfig
>>> config = LDMemoryConfig(memory_limit_gb=8.0)
>>> result = fit(model, xdata, ydata, p0=[1, 2], workflow=config)

See also

curve_fit

Lower-level API with full control

curve_fit_large

Specialized API for large datasets

MemoryBudgetSelector

Memory-based automatic strategy selection