1.2. Optimization Pipeline

This page explains how data flows through NLSQ’s optimization pipeline.

1.2.1. Pipeline Overview

User Code
    │
    ▼
fit() / curve_fit()              ◄── Entry point
    │
    ▼
┌─────────────────────────────────────────────────┐
│                  CurveFit                        │
│  ┌─────────────────────────────────────────────┐│
│  │ DataPreprocessor                             ││
│  │ - Validate inputs                           ││
│  │ - Convert to JAX arrays                     ││
│  │ - Handle NaN, masking                       ││
│  └─────────────────────────────────────────────┘│
│                    │                             │
│                    ▼                             │
│  ┌─────────────────────────────────────────────┐│
│  │ OptimizationSelector                         ││
│  │ - Detect parameter count                    ││
│  │ - Process bounds                            ││
│  │ - Generate initial guess                    ││
│  │ - Select method/solver                      ││
│  └─────────────────────────────────────────────┘│
│                    │                             │
│                    ▼                             │
│  ┌─────────────────────────────────────────────┐│
│  │ StreamingCoordinator                         ││
│  │ - Analyze memory requirements               ││
│  │ - Select: STANDARD / CHUNKED / STREAMING    ││
│  └─────────────────────────────────────────────┘│
└────────────────────┬────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────┐
│              LeastSquares                        │
│  - Create residual function                     │
│  - Configure Jacobian (auto-diff)               │
│  - Set up trust region options                  │
└────────────────────┬────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────┐
│          TrustRegionReflective                   │
│  - Iterative optimization loop                  │
│  - Trust region updates                         │
│  - Convergence checking                         │
└────────────────────┬────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────┐
│           CovarianceComputer                     │
│  - Compute Jacobian at solution                 │
│  - SVD for covariance estimation                │
│  - Apply sigma transformation                   │
└────────────────────┬────────────────────────────┘
                     │
                     ▼
                (popt, pcov)

1.2.2. Step-by-Step Walkthrough

1. Entry Point (fit/curve_fit)

from nlsq import fit

popt, pcov = fit(model, x, y, p0=[1, 0.5])
# fit() creates CurveFit instance and calls curve_fit()

2. Data Preprocessing

# DataPreprocessor handles:
# - Type conversion to JAX arrays
# - NaN handling based on nan_policy
# - Data masking
# - Shape validation

preprocessed = DataPreprocessor().preprocess(
    f=model, xdata=x, ydata=y, sigma=None, nan_policy="raise"
)

3. Optimization Selection

# OptimizationSelector determines:
# - Number of parameters (from model signature)
# - Bounds processing
# - Initial guess validation
# - Method selection (trf, dogbox)

config = OptimizationSelector().select(
    f=model, xdata=x, ydata=y, p0=[1, 0.5], bounds=None, method="trf"
)

4. Memory Strategy Selection

# StreamingCoordinator analyzes memory:
# - Data size vs available memory
# - Peak memory estimation
# - Strategy selection

decision = StreamingCoordinator().decide(xdata=x, ydata=y, n_params=2, workflow="auto")
# Returns: strategy='standard', 'chunked', or 'streaming'

5. LeastSquares Optimization

# LeastSquares sets up the optimization:
# - Wraps model as residual function
# - Configures automatic differentiation
# - Calls TrustRegionReflective

optimizer = LeastSquares()
result = optimizer.least_squares(
    fun=residuals,
    x0=p0,
    jac="2-point",  # or analytical
    bounds=(-np.inf, np.inf),
    method="trf",
)

6. Trust Region Iteration

# TRF performs iterative optimization:
while not converged:
    # 1. Compute Jacobian (auto-diff or finite difference)
    J = compute_jacobian(x_current)

    # 2. Solve trust region subproblem
    step = solve_subproblem(J, residuals, trust_radius)

    # 3. Evaluate new point
    x_new = x_current + step
    cost_new = compute_cost(x_new)

    # 4. Update trust region
    if cost_new < cost_current:
        x_current = x_new
        expand_trust_region()
    else:
        contract_trust_region()

    # 5. Check convergence
    converged = check_convergence(ftol, xtol, gtol)

7. Covariance Computation

# CovarianceComputer estimates uncertainties:
# - Final Jacobian at solution
# - SVD decomposition
# - Covariance from inverse of J^T J

cov_result = CovarianceComputer().compute(
    result=optimize_result, n_data=len(y), sigma=None, absolute_sigma=False
)

1.2.3. Jacobian Computation

NLSQ uses JAX autodiff for Jacobians:

import jax

# Forward-mode: efficient for few parameters
J = jax.jacfwd(residual_func)(params)

# Reverse-mode: efficient for many parameters
J = jax.jacrev(residual_func)(params)

# Auto-selection based on dimensions
# n_params < n_residuals → forward-mode
# n_params > n_residuals → reverse-mode

1.2.4. Global Optimization Path

For workflow='auto_global':

fit(workflow='auto_global')
        │
        ▼
MethodSelector
- scale_ratio > 1000? → CMA-ES
- otherwise → Multi-Start
        │
        ├──► MultiStartOrchestrator
        │       │
        │       ├── Latin Hypercube Sampling
        │       ├── n parallel local optimizations
        │       └── Select best result
        │
        └──► CMAESOptimizer (if selected)
                │
                ├── Evolutionary search
                ├── BIPOP restarts
                └── Local refinement

1.2.5. Next Steps