Core API Reference

This page documents the main functions and classes for curve fitting.

Primary Functions

fit()

nlsq.fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False, check_finite=True, bounds=(-inf, inf), method=None, workflow=None, preset=None, multistart=None, n_starts=None, sampler='lhs', center_on_p0=True, scale_factor=1.0, memory_limit_gb=None, size_threshold=1000000, show_progress=False, chunk_size=None, **kwargs)[source]

Unified curve fitting function with preset-based configuration.

This function provides a simplified API for curve fitting with sensible defaults based on preset configurations. It automatically selects the appropriate backend (curve_fit, curve_fit_large, or streaming) based on the preset and dataset characteristics.

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).

  • preset ({'fast', 'robust', 'global', 'streaming', 'large'}, optional) –

    Preset configuration to use:

    • ’fast’: Single-start optimization for maximum speed (n_starts=0)

    • ’robust’: Multi-start with 5 starts for robustness

    • ’global’: Thorough global search with 20 starts

    • ’streaming’: Streaming optimization for large datasets with multi-start

    • ’large’: Auto-detect dataset size and use appropriate strategy

    If None, defaults to ‘fast’ for small datasets or ‘large’ for datasets exceeding size_threshold.

  • multistart (bool, optional) – Override preset’s multi-start setting.

  • n_starts (int, optional) – Override preset’s n_starts setting.

  • sampler ({'lhs', 'sobol', 'halton'}, optional) – Sampling strategy for multi-start. Default: ‘lhs’.

  • center_on_p0 (bool, optional) – Center multi-start samples around p0. Default: True.

  • scale_factor (float, optional) – Scale factor for exploration region. Default: 1.0.

  • memory_limit_gb (float, optional) – Maximum memory usage in GB for large datasets.

  • size_threshold (int, optional) – Point threshold for large dataset processing (default: 1M).

  • show_progress (bool, optional) – Display progress bar for long operations.

  • chunk_size (int, optional) – Override automatic chunk size calculation.

  • **kwargs (Any) – Additional optimization parameters (ftol, xtol, gtol, max_nfev, loss).

Returns:

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

Return type:

CurveFitResult or tuple

Examples

Basic usage with default preset:

>>> popt, pcov = fit(model_func, xdata, ydata, p0=[1, 2, 3])

Using ‘robust’ preset for multi-start:

>>> result = fit(model_func, xdata, ydata, p0=[1, 2, 3],
...              bounds=([0, 0, 0], [10, 10, 10]), preset='robust')

Using ‘global’ preset for thorough search:

>>> result = fit(model_func, xdata, ydata, p0=[1, 2, 3],
...              bounds=([0, 0, 0], [10, 10, 10]), preset='global')

Large dataset with auto-detection:

>>> result = fit(model_func, big_xdata, big_ydata,
...              preset='large', show_progress=True)

See also

curve_fit

Lower-level API with full control

curve_fit_large

Specialized API for large datasets

MemoryBudgetSelector

Memory-based optimizer selection

The fit() function is the recommended entry point. It provides a unified interface with preset-based configuration.

Presets:

Preset

Description

"default"

Standard configuration for most use cases

"fast"

Optimized for speed, lower precision

"precise"

Higher precision, more iterations

"global"

Global optimization to avoid local minima

"streaming"

For datasets that exceed memory

Example:

from nlsq import fit
import jax.numpy as jnp


def model(x, a, b, c):
    return a * jnp.exp(-b * x) + c


popt, pcov = fit(model, x_data, y_data, p0=[1.0, 0.5, 0.0])

curve_fit()

nlsq.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
... )

Drop-in replacement for scipy.optimize.curve_fit. Use this when migrating from SciPy or when you need precise control over optimization parameters.

Key Parameters:

  • f: Model function with signature f(x, *params)

  • xdata: Independent variable data

  • ydata: Dependent variable data

  • p0: Initial parameter guesses

  • sigma: Measurement uncertainties

  • bounds: Parameter bounds as ([lowers], [uppers])

  • max_nfev: Maximum function evaluations

  • gtol, ftol, xtol: Convergence tolerances

Example:

from nlsq import curve_fit
import jax.numpy as jnp


def exponential(x, a, tau, c):
    return a * jnp.exp(-x / tau) + c


popt, pcov = curve_fit(
    exponential, x_data, y_data, p0=[1.0, 10.0, 0.0], bounds=([0, 0, -1], [10, 100, 1])
)

curve_fit_large()

nlsq.curve_fit_large(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False, check_finite=True, bounds=(-inf, inf), method=None, stability=False, rescale_data=True, max_jacobian_elements_for_svd=10000000, memory_limit_gb=None, auto_size_detection=True, size_threshold=1000000, show_progress=False, chunk_size=None, multistart=False, n_starts=10, global_search=False, sampler='lhs', center_on_p0=True, scale_factor=1.0, **kwargs)[source]

Curve fitting with automatic memory management for large datasets.

Automatically selects processing strategy based on dataset size: - Small (< 1M points): Standard curve_fit - Medium (1M - 100M points): Chunked processing - Large (> 100M points): Streaming optimization

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).

  • memory_limit_gb (float, optional) – Maximum memory usage in GB.

  • auto_size_detection (bool, optional) – Auto-detect dataset size for processing strategy.

  • size_threshold (int, optional) – Point threshold for large dataset processing (default: 1M).

  • show_progress (bool, optional) – Display progress bar for long operations.

  • chunk_size (int, optional) – Override automatic chunk size calculation.

  • multistart (bool, optional) – Enable multi-start optimization for global search. Default: False.

  • n_starts (int, optional) – Number of starting points for multi-start optimization. Default: 10.

  • global_search (bool, optional) – Shorthand for multistart=True, n_starts=20. Default: False.

  • sampler ({'lhs', 'sobol', 'halton'}, optional) – Sampling strategy for multi-start. Default: ‘lhs’.

  • center_on_p0 (bool, optional) – Center multi-start samples around p0. Default: True.

  • scale_factor (float, optional) – Scale factor for exploration region. Default: 1.0.

  • **kwargs (Any) – Additional optimization parameters (ftol, xtol, gtol, max_nfev, loss)

Returns:

  • popt (ndarray) – Fitted parameters.

  • pcov (ndarray) – Parameter covariance matrix.

Return type:

tuple[ndarray, ndarray] | OptimizeResult

Notes

All large datasets use streaming optimization for 100% data utilization.

Important: Model Function Requirements for Chunking

When auto_size_detection triggers chunked processing (>1M points), your model function MUST respect the size of xdata. Model output shape must match ydata shape.

INCORRECT - Fixed-size output (causes shape errors):

>>> def bad_model(xdata, a, b):
...     # WRONG: Returns fixed-size array regardless of xdata
...     t_full = jnp.arange(10_000_000)
...     return a * jnp.exp(-b * t_full)  # Shape mismatch!

CORRECT - Output matches xdata size:

>>> def good_model(xdata, a, b):
...     # CORRECT: Uses xdata as indices
...     indices = xdata.astype(jnp.int32)
...     return a * jnp.exp(-b * indices)

>>> def direct_model(xdata, a, b):
...     # CORRECT: Operates on xdata directly
...     return a * jnp.exp(-b * xdata)

Examples

Basic usage:

>>> popt, _pcov = curve_fit_large(model_func, xdata, ydata, p0=[1, 2, 3])

Large dataset with progress bar:

>>> popt, _pcov = curve_fit_large(model_func, big_xdata, big_ydata,
...                             show_progress=True, memory_limit_gb=8)

With multi-start optimization:

>>> popt, _pcov = curve_fit_large(model_func, xdata, ydata,
...                             p0=[1, 2, 3], bounds=([0, 0, 0], [10, 10, 10]),
...                             multistart=True, n_starts=10)

Using external logger for diagnostics:

>>> import logging
>>> my_logger = logging.getLogger("myapp")
>>> fitter = LargeDatasetFitter(memory_limit_gb=8, logger=my_logger)
>>> result = fitter.fit(model_func, xdata, ydata, p0=[1, 2])
>>> # Chunk failures now appear in myapp's logs

Handles datasets that exceed GPU memory through automatic chunking.

Key Parameters:

  • memory_limit_gb: Maximum GPU memory to use (default: auto-detect)

  • chunk_size: Manual chunk size (overrides auto)

  • All parameters from curve_fit()

Example:

from nlsq import curve_fit_large

# Fit 10 million points with automatic chunking
popt, pcov = curve_fit_large(model, x_large, y_large, p0=p0, memory_limit_gb=8.0)

LeastSquares

class nlsq.LeastSquares(enable_stability=False, enable_diagnostics=False, max_jacobian_elements_for_svd=10000000)[source]

Bases: object

Core least squares optimization engine with JAX acceleration.

This class implements the main optimization algorithms for nonlinear least squares problems, including Trust Region Reflective (TRF) and Levenberg-Marquardt (LM). It handles automatic differentiation, bound constraints, loss functions, and uncertainty propagation.

The class maintains separate automatic differentiation instances for different sigma configurations (no sigma, 1D sigma, 2D covariance matrix) to optimize compilation and execution performance.

trf

Trust Region Reflective algorithm implementation

Type:

TrustRegionReflective

ls

JIT-compiled loss function implementations

Type:

LossFunctionsJIT

logger

Internal logger for debugging and performance tracking

Type:

Logger

f

Current objective function being optimized

Type:

callable

jac

Current Jacobian function (None for automatic differentiation)

Type:

callable or None

adjn

Automatic differentiation instance for unweighted problems

Type:

AutoDiffJacobian

adj1d

Automatic differentiation instance for 1D sigma weighting

Type:

AutoDiffJacobian

adj2d

Automatic differentiation instance for 2D covariance matrix weighting

Type:

AutoDiffJacobian

least_squares : Main optimization method
__init__(enable_stability=False, enable_diagnostics=False, max_jacobian_elements_for_svd=10000000)[source]

Initialize LeastSquares with optimization algorithms and autodiff instances.

Sets up the Trust Region Reflective solver, loss functions, and separate automatic differentiation instances for different weighting schemes to maximize JAX compilation efficiency.

Parameters:

  • enable_stability (bool, default False) – Enable numerical stability checks and fixes

  • enable_diagnostics (bool, default False) – Enable optimization diagnostics collection

  • max_jacobian_elements_for_svd (int, default 10_000_000) – Maximum Jacobian size (m × n elements) for SVD computation during stability checks. SVD is skipped for larger Jacobians.

least_squares(fun, x0, jac=None, bounds=(-inf, inf), method='trf', ftol=1e-08, xtol=1e-08, gtol=1e-08, x_scale=1.0, loss='linear', f_scale=1.0, diff_step=None, tr_solver=None, tr_options=None, jac_sparsity=None, max_nfev=None, verbose=0, jacobian_mode=None, xdata=None, ydata=None, data_mask=None, transform=None, timeit=False, callback=None, args=(), kwargs=None, prepared_bounds=None, **timeout_kwargs)[source]

Solve nonlinear least squares problem using JAX-accelerated algorithms.

This method orchestrates the optimization process by calling focused helper methods for each major step: validation, function setup, initial evaluation, stability checks, and optimization execution.

Parameters:

  • fun (callable) – Residual function. Must use jax.numpy operations.

  • x0 (array_like) – Initial parameter guess.

  • jac (callable or None, optional) – Jacobian function. If None, uses JAX autodiff.

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

  • method (str, optional) – Optimization algorithm (‘trf’).

  • ftol (float, optional) – Convergence tolerances for function, parameters, and gradient.

  • xtol (float, optional) – Convergence tolerances for function, parameters, and gradient.

  • gtol (float, optional) – Convergence tolerances for function, parameters, and gradient.

  • x_scale (str or array_like, optional) – Parameter scaling (‘jac’ for automatic).

  • loss (str or callable, optional) – Robust loss function (‘linear’, ‘huber’, ‘soft_l1’, etc.).

  • f_scale (float, optional) – Scale parameter for robust loss functions.

  • max_nfev (int, optional) – Maximum function evaluations.

  • verbose (int, optional) – Verbosity level (0, 1, or 2).

  • jacobian_mode ({'auto', 'fwd', 'rev'}, optional) – Jacobian automatic differentiation mode. If None, uses configuration from environment variable, config file, or auto-default. Default is None. - ‘auto’: Automatically select based on problem dimensions - ‘fwd’: Force forward-mode AD (jacfwd) - ‘rev’: Force reverse-mode AD (jacrev)

  • xdata (array_like, optional) – Data for curve fitting applications.

  • ydata (array_like, optional) – Data for curve fitting applications.

  • data_mask (array_like, optional) – Boolean mask for data exclusion.

  • transform (array_like, optional) – Transformation matrix for weighted fitting.

  • timeit (bool, optional) – Enable detailed timing analysis.

  • callback (callable or None, optional) – Callback function called after each optimization iteration with signature callback(iteration, cost, params, info). Useful for monitoring optimization progress, logging, or implementing custom stopping criteria. If None (default), no callback is invoked.

  • args (tuple, optional) – Additional arguments for objective function.

  • kwargs (dict, optional) – Additional optimization parameters.

  • prepared_bounds (tuple[np.ndarray, np.ndarray] or None, optional) – Pre-prepared bounds as (lb, ub). If provided, skips prepare_bounds call. This optimization avoids redundant bounds preparation when caller has already called prepare_bounds. Default is None.

Returns:

result – Optimization result with solution, convergence info, and statistics.

Return type:

OptimizeResult

autdiff_jac(jac, mode='fwd')[source]

We do this for all three sigma transformed functions such that if sigma is changed from none to 1D to covariance sigma then no retracing is needed.

Parameters:

  • jac (None) – Passed in to maintain compatibility with the user defined Jacobian function.

  • mode (str, optional) – Jacobian mode (‘fwd’ or ‘rev’), by default ‘fwd’

update_function(func)[source]

Wraps the given fit function to be a residual function using the data. The wrapped function is in a JAX JIT compatible format which is purely functional. This requires that both the data mask and the uncertainty transform are passed to the function. Even for the case where the data mask is all True and the uncertainty transform is None we still need to pass these arguments to the function due JAX’s functional nature.

Parameters:

func (Callable) – The fit function to wrap.

Return type:

None

wrap_jac(jac)[source]

Wraps an user defined Jacobian function to allow for data masking and uncertainty transforms. The wrapped function is in a JAX JIT compatible format which is purely functional. This requires that both the data mask and the uncertainty transform are passed to the function.

Using an analytical Jacobian of the fit function is equivalent to the Jacobian of the residual function.

Also note that the analytical Jacobian doesn’t require the independent ydata, but we still need to pass it to the function to maintain compatibility with autdiff version which does require the ydata.

Parameters:

jac (Callable) – The Jacobian function to wrap.

Returns:

The masked Jacobian of the function evaluated at args with respect to the data.

Return type:

jnp.ndarray

Low-level least squares solver class. Use this when you need full control over the optimization process or are working with non-standard objective functions.

Classes

CurveFit

class nlsq.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’.

Reusable curve fitting class that caches JIT compilation for repeated fits.

Example:

from nlsq import CurveFit

fitter = CurveFit()

# First call compiles (slower)
result1 = fitter.curve_fit(model, x1, y1, p0=p0)

# Subsequent calls use cached compilation (fast)
result2 = fitter.curve_fit(model, x2, y2, p0=p0)
result3 = fitter.curve_fit(model, x3, y3, p0=p0)

OptimizeResult

class nlsq.result.OptimizeResult[source]

Bases: dict

Optimization result container for NLSQ curve fitting operations.

This class stores the complete results from nonlinear least squares optimization performed using JAX-accelerated algorithms. It extends dict to provide both dictionary-style and attribute-style access to optimization results.

Core Attributes

xjax.numpy.ndarray or numpy.ndarray

Optimized parameter vector containing the final fitted parameters. These represent the solution to the nonlinear least squares problem.

successbool

Indicates whether the optimization terminated successfully. True means convergence criteria were satisfied within tolerance limits.

statusint

Numerical termination status code indicating why optimization stopped:

  • 1: Gradient convergence (||g||_inf < gtol)

  • 2: Step size convergence (||dx||/||x|| < xtol)

  • 3: Function value convergence (delta_f/f < ftol)

  • 0: Maximum iterations reached

  • -1: Evaluation limit exceeded

  • -3: Inner loop iteration limit (algorithm-specific)

messagestr

Human-readable description of termination cause. Provides detailed information about convergence status or failure reasons.

Objective Function Results

funjax.numpy.ndarray

Final residual vector f(x) at the solution. For curve fitting, these are the differences between model predictions and data points.

costfloat

Final cost function value: 0.5 * ||f(x)||² for standard least squares, or 0.5 * sum(ρ(f_i²/σ²)) for robust loss functions.

jacjax.numpy.ndarray

Final Jacobian matrix J(x) with shape (m, n) where m is number of data points and n is number of parameters. Computed using JAX autodiff.

gradjax.numpy.ndarray

Final gradient vector g = J^T * f with shape (n,). Used for convergence checking and parameter uncertainty estimation.

Convergence Metrics

optimalityfloat

Final gradient norm ||g||_inf used for convergence assessment. Should be less than gtol for successful convergence.

active_masknumpy.ndarray

Boolean mask indicating which parameters hit bounds (for bounded optimization). Shape (n,) with True for parameters at constraints.

Iteration Statistics

nfevint

Total number of objective function evaluations during optimization. Each evaluation computes residuals f(x) for given parameters.

njevint

Total number of Jacobian evaluations. With JAX autodiff, this equals the number of combined function+gradient evaluations.

nitint

Number of optimization iterations completed. Not always available for all algorithms.

Algorithm-Specific Results

pcovjax.numpy.ndarray, optional

Parameter covariance matrix with shape (n, n). Provides parameter uncertainty estimates. Available when uncertainty estimation is requested. Computed as: pcov = inv(J^T * J) * residual_variance

active_masknumpy.ndarray

For bounded optimization, indicates which parameters are at bounds.

all_timesdict, optional

Detailed timing information for algorithm profiling. Contains timing data for different optimization phases (function evaluation, Jacobian computation, linear algebra operations, etc.).

Usage Examples

Basic result access:

import nlsq

# Perform curve fitting
result = nlsq.curve_fit(model_func, x_data, y_data, p0=initial_guess)

# Access optimized parameters
fitted_params = result.x

# Check convergence
if result.success:
    print(f"Optimization converged: {result.message}")
    print(f"Final cost: {result.cost}")
    print(f"Function evaluations: {result.nfev}")
else:
    print(f"Optimization failed: {result.message}")

# Parameter uncertainties (if covariance computed)
if hasattr(result, 'pcov'):
    param_errors = jnp.sqrt(jnp.diag(result.pcov))
    print(f"Parameter uncertainties: {param_errors}")

Advanced result inspection:

# Examine residuals and fit quality
final_residuals = result.fun
rms_error = jnp.sqrt(jnp.mean(final_residuals**2))

# Check gradient convergence
gradient_norm = result.optimality
print(f"Final gradient norm: {gradient_norm}")

# Analyze Jacobian condition
jacobian = result.jac
condition_number = jnp.linalg.cond(jacobian)
print(f"Jacobian condition number: {condition_number}")

# For bounded problems, check active constraints
if hasattr(result, 'active_mask'):
    constrained_params = jnp.where(result.active_mask)[0]
    print(f"Parameters at bounds: {constrained_params}")

Integration with SciPy

This class maintains compatibility with scipy.optimize.OptimizeResult while adding JAX-specific features and NLSQ-specific results. It can be used interchangeably with SciPy optimization results in most contexts.

Technical Notes

  • All JAX arrays are automatically converted to NumPy arrays for compatibility

  • Covariance matrices use double precision for numerical stability

  • Large dataset results may include memory management statistics

  • GPU timing results require explicit timing mode activation

  • Progress monitoring data is stored in algorithm-specific attributes

The result object returned by optimization functions.

Key Attributes:

  • x: Optimal parameters (same as popt in curve_fit)

  • cost: Final value of the cost function

  • fun: Final residuals

  • jac: Final Jacobian matrix

  • nfev: Number of function evaluations

  • nit: Number of iterations

  • message: Termination message

  • success: Whether optimization succeeded

Methods:

  • plot(): Plot data, fit, and confidence bands

  • confidence_band(): Calculate confidence intervals

See Visualization API for details on plotting.

Return Values

curve_fit Return Format

popt, pcov = curve_fit(model, x, y, p0=p0)

# popt: Optimal parameter values (JAX array)
# pcov: Covariance matrix of parameters (JAX array)

# Standard errors
perr = jnp.sqrt(jnp.diag(pcov))

fit Return Format

popt, pcov = fit(model, x, y, p0=p0)

# Same as curve_fit

# With full_output=True:
popt, pcov, info = fit(model, x, y, p0=p0, full_output=True)

# info contains:
# - residuals: Final residuals
# - jacobian: Final Jacobian
# - nfev: Function evaluations
# - nit: Iterations

Error Handling

NLSQ raises informative exceptions:

from nlsq import curve_fit
from nlsq.exceptions import ConvergenceError, ValidationError

try:
    popt, pcov = curve_fit(model, x, y, p0=p0)
except ConvergenceError as e:
    print(f"Did not converge: {e}")
except ValidationError as e:
    print(f"Invalid input: {e}")

See Also