Curve Fitting

Overview

The piblin_jax.fitting module provides non-linear least squares (NLSQ) curve fitting functionality for piblin_jax. It offers a fast, deterministic approach to parameter estimation when uncertainty quantification is not required.

This module integrates with the JAX-based NLSQ library for high-performance optimization on CPU/GPU, with automatic fallback to scipy’s curve_fit when JAX is unavailable. The fitting functions are designed to work seamlessly with piblin-jax’s data structures while providing a simple, familiar interface.

Key characteristics:

Performance: JAX-based optimization with JIT compilation provides significant speed improvements over traditional scipy methods, especially for complex models or large datasets. GPU acceleration is automatically utilized when available.
Automatic Initial Estimates: The module includes heuristics for automatically estimating initial parameter values for common model types, reducing the need for manual tuning.
Robust Optimization: The underlying NLSQ implementation uses the Levenberg-Marquardt algorithm, which balances the speed of gradient descent with the stability of Gauss-Newton methods.
Scipy Compatibility: When JAX is not available, the module automatically falls back to scipy.optimize.curve_fit, ensuring broad compatibility across different environments.

When to use fitting vs. Bayesian models:

Use fitting when you need fast parameter estimates, have clean data, and don’t need uncertainty quantification. Ideal for exploratory analysis, real-time processing, or when computational resources are limited.
Use Bayesian models when you need full uncertainty quantification, want to incorporate prior knowledge, have noisy/limited data, or need to compare competing models rigorously.

Quick Examples

Basic Curve Fitting

Fit a custom function to data:

from piblin_jax.fitting import fit_curve
import numpy as np

def power_law(x, K, n):
    """Power law model: y = K * x^n"""
    return K * x**n

# Prepare data
x_data = np.logspace(-2, 3, 50)
y_data = np.array([...])  # Your measurements

# Fit the model
params, covariance = fit_curve(
    f=power_law,
    xdata=x_data,
    ydata=y_data,
    p0=[1.0, 0.5]  # Initial guess: K=1.0, n=0.5
)

K_opt, n_opt = params
print(f"Fitted parameters: K={K_opt:.3f}, n={n_opt:.3f}")

Automatic Initial Parameter Estimation

Let the module estimate initial parameters for you:

from piblin_jax.fitting import fit_curve, estimate_initial_parameters

# Estimate initial parameters automatically
p0 = estimate_initial_parameters(
    f=power_law,
    xdata=x_data,
    ydata=y_data,
    model_type="power_law"  # Hint for better estimates
)

# Fit with estimated parameters
params, covariance = fit_curve(
    f=power_law,
    xdata=x_data,
    ydata=y_data,
    p0=p0
)

Working with Bounds and Constraints

Apply parameter bounds for physically meaningful results:

# Fit with parameter bounds
params, covariance = fit_curve(
    f=power_law,
    xdata=x_data,
    ydata=y_data,
    p0=[1.0, 0.5],
    bounds=([0.0, 0.0], [np.inf, 2.0])  # K > 0, 0 < n < 2
)

# Extract parameter uncertainties from covariance
param_std = np.sqrt(np.diag(covariance))
print(f"K = {params[0]:.3f} +/- {param_std[0]:.3f}")
print(f"n = {params[1]:.3f} +/- {param_std[1]:.3f}")

API Reference

Module Contents

Curve fitting module for piblin-jax.

This module provides non-linear least squares (NLSQ) curve fitting for rheological and scientific models, with automatic parameter initialization and robust optimization.

## Overview

The fitting module bridges piblin-jax’s Bayesian models with classical NLSQ optimization, providing: - Fast parameter estimation for model initialization - Multiple rheological models (power-law, Arrhenius, Cross, Carreau-Yasuda) - Automatic initial parameter guessing via heuristics - scipy.optimize integration with robust error handling - Parameter uncertainty estimates via covariance matrix

## When to Use

NLSQ fitting: Quick parameter estimates, no uncertainty quantification
Bayesian fitting (piblin_jax.bayesian): Full posterior distributions, uncertainty propagation

Typical workflow: 1. Use fit_curve() for initial parameter estimates 2. Use Bayesian models for comprehensive uncertainty quantification 3. Or use NLSQ estimates to initialize Bayesian priors

## Main Functions

### fit_curve()

Fit rheological models to experimental data using non-linear least squares:

from piblin_jax import fit_curve
import numpy as np

shear_rate = np.logspace(-1, 2, 30)
viscosity = 5.0 * shear_rate ** (0.6 - 1)

result = fit_curve(shear_rate, viscosity, model='power_law')
print(result['params'])  # {'K': 5.02, 'n': 0.598}

Supported Models: - ‘power_law’: η = K * γ̇^(n-1) - ‘arrhenius’: η = A * exp(Ea / (R*T)) - ‘cross’: η = η∞ + (η₀ - η∞) / (1 + (λγ̇)^m) - ‘carreau_yasuda’: η = η∞ + (η₀ - η∞) * [1 + (λγ̇)^a]^((n-1)/a)

### estimate_initial_parameters()

Automatically estimate initial parameter guesses for optimization:

from piblin_jax import estimate_initial_parameters

initial = estimate_initial_parameters(
    shear_rate,
    viscosity,
    model='power_law'
)
print(initial)  # {'K': 4.8, 'n': 0.55}

Methods: - Power-law: Linear regression on log-log data - Arrhenius: Linear regression on log(η) vs 1/T - Cross/Carreau-Yasuda: Plateau detection + heuristics

## Examples

### Basic Fitting

Example:

import numpy as np
from piblin_jax import fit_curve

# Generate data
shear_rate = np.logspace(-2, 3, 50)
viscosity = 5.0 * shear_rate ** (0.6 - 1)

# Fit model
result = fit_curve(shear_rate, viscosity, model='power_law')

# Extract results
params = result['params']
K, n = params['K'], params['n']
covariance = result['covariance']
residuals = result['residuals']

print(f"K = {K:.3f} ± {np.sqrt(covariance[0,0]):.3f}")
print(f"n = {n:.3f} ± {np.sqrt(covariance[1,1]):.3f}")

### With Initial Guesses

Example:

# Provide initial guesses
initial = {'K': 3.0, 'n': 0.5}
result = fit_curve(
    shear_rate,
    viscosity,
    model='power_law',
    initial_params=initial
)

### Multiple Models

Example:

models = ['power_law', 'cross', 'carreau_yasuda']
results = {}

for model in models:
    try:
        results[model] = fit_curve(shear_rate, viscosity, model=model)
        print(f"{model}: RSS = {np.sum(results[model]['residuals']**2):.3e}")
    except Exception as e:
        print(f"{model} failed: {e}")

### NLSQ → Bayesian Workflow

Example:

from piblin_jax import fit_curve
from piblin_jax.bayesian.models import PowerLawModel

# Step 1: Quick NLSQ estimate
nlsq_result = fit_curve(shear_rate, viscosity, model='power_law')
print(f"NLSQ estimate: K={nlsq_result['params']['K']:.2f}")

# Step 2: Bayesian fitting with uncertainty
model = PowerLawModel(n_samples=2000)
model.fit(shear_rate, viscosity)

# Step 3: Compare results
bayesian_K = np.mean(model.samples['K'])
K_std = np.std(model.samples['K'])
print(f"Bayesian estimate: K={bayesian_K:.2f} ± {K_std:.2f}")

## Return Values

fit_curve() returns a dictionary with: - ‘params’: dict[str, float] - Fitted parameter values - ‘covariance’: np.ndarray - Parameter covariance matrix - ‘residuals’: np.ndarray - Fitting residuals (y_data - y_fit) - ‘success’: bool - Optimization convergence status - ‘message’: str - Optimization status message - ‘nfev’: int - Number of function evaluations

## Error Handling

Example:

try:
    result = fit_curve(x, y, model='power_law')
    if not result['success']:
        print(f"Optimization warning: {result['message']}")
        # Still can use params, but check residuals
except ValueError as e:
    print(f"Invalid input: {e}")
except RuntimeError as e:
    print(f"Optimization failed: {e}")

## Implementation Details

Optimization: - Backend: scipy.optimize.curve_fit with Levenberg-Marquardt algorithm - Automatic bounds: Parameters constrained to physical ranges - Robust initialization: Automatic parameter guessing prevents convergence failures

Performance: - Typical fit time: 1-10 ms (CPU) - No GPU acceleration (classical optimization) - For GPU speedup, use Bayesian models with JAX backend

## See Also

piblin_jax.bayesian.models - Bayesian fitting with uncertainty quantification
piblin_jax.bayesian.models.PowerLawModel - Bayesian power-law fitting
piblin_jax.bayesian.models.ArrheniusModel - Bayesian Arrhenius fitting
scipy.optimize.curve_fit - Underlying optimization routine

## References

[1]

More, J.J., Garbow, B.S., and Hillstrom, K.E. (1980). “User Guide for MINPACK-1”, Argonne National Laboratory.

[2]

Marquardt, D.W. (1963). “An Algorithm for Least-Squares Estimation of Nonlinear Parameters”, SIAM Journal on Applied Mathematics.

piblin_jax.fitting.estimate_initial_parameters(func, x, y, bounds=None)[source]

Estimate initial parameters for curve fitting.

Uses simple heuristics to estimate reasonable initial parameter values based on the data characteristics.

Parameters:

func (Callable) – Model function (used to determine number of parameters).
x (np.ndarray) – Independent variable data.
y (np.ndarray) – Dependent variable data.
bounds (tuple, optional) – Parameter bounds (lower, upper).

Returns:

Estimated initial parameters.

Return type:

np.ndarray

Notes

This is a simple heuristic approach. For better results, use domain knowledge to provide good initial guesses.

piblin_jax.fitting.fit_curve(func, x, y, p0=None, sigma=None, absolute_sigma=False, **kwargs)[source]

Fit a curve using NLSQ or scipy fallback.

This function attempts to use the NLSQ library for curve fitting. If NLSQ is not available, it falls back to scipy.optimize.curve_fit.

Parameters:

func (Callable) – Model function to fit. Should have signature func(x, *params).
x (np.ndarray) – Independent variable data.
y (np.ndarray) – Dependent variable data.
p0 (np.ndarray, optional) – Initial parameter guess. If None, will use default initialization.
sigma (np.ndarray, optional) – Uncertainty/weights for data points. If provided, used in weighted fit.
absolute_sigma (bool, default False) – If True, sigma is used in absolute sense and parameter covariance matrix is scaled accordingly.
**kwargs (Any) – Additional keyword arguments passed to fitting function.

Returns:

Dictionary containing: - ‘params’: Fitted parameters - ‘covariance’: Parameter covariance matrix - ‘method’: Method used (‘nlsq’ or ‘scipy’) - ‘success’: Whether fit converged - ‘residuals’: Residuals (y - func(x, *params))

Return type:

dict

Examples

>>> import numpy as np
>>> from piblin_jax.fitting.nlsq import fit_curve
>>>
>>> # Define model function
>>> def linear_model(x, a, b):
...     return a * x + b
>>>
>>> # Generate data
>>> x = np.linspace(0, 10, 50)
>>> y = 2.5 * x + 1.0 + 0.1 * np.random.randn(len(x))
>>>
>>> # Fit
>>> result = fit_curve(linear_model, x, y, p0=[1.0, 0.0])
>>> print(f"Fitted parameters: {result['params']}")
>>> print(f"Method used: {result['method']}")

Notes

NLSQ often provides better convergence than scipy for challenging problems
The scipy fallback ensures the function always works
For Bayesian inference, consider using BayesianModel instead

Non-Linear Least Squares

NLSQ integration for curve fitting.

This module provides a wrapper around the NLSQ library for nonlinear least squares curve fitting, with automatic fallback to scipy.optimize.curve_fit if NLSQ is not available.

The NLSQ library (if available) provides enhanced nonlinear least squares fitting with better convergence properties than scipy for many problems.

piblin_jax.fitting.nlsq.estimate_initial_parameters(func, x, y, bounds=None)[source]

Estimate initial parameters for curve fitting.

Uses simple heuristics to estimate reasonable initial parameter values based on the data characteristics.

Parameters:

func (Callable) – Model function (used to determine number of parameters).
x (np.ndarray) – Independent variable data.
y (np.ndarray) – Dependent variable data.
bounds (tuple, optional) – Parameter bounds (lower, upper).

Returns:

Estimated initial parameters.

Return type:

np.ndarray

Notes

This is a simple heuristic approach. For better results, use domain knowledge to provide good initial guesses.

piblin_jax.fitting.nlsq.fit_curve(func, x, y, p0=None, sigma=None, absolute_sigma=False, **kwargs)[source]

Fit a curve using NLSQ or scipy fallback.

This function attempts to use the NLSQ library for curve fitting. If NLSQ is not available, it falls back to scipy.optimize.curve_fit.

Parameters:

func (Callable) – Model function to fit. Should have signature func(x, *params).
x (np.ndarray) – Independent variable data.
y (np.ndarray) – Dependent variable data.
p0 (np.ndarray, optional) – Initial parameter guess. If None, will use default initialization.
sigma (np.ndarray, optional) – Uncertainty/weights for data points. If provided, used in weighted fit.
absolute_sigma (bool, default False) – If True, sigma is used in absolute sense and parameter covariance matrix is scaled accordingly.
**kwargs (Any) – Additional keyword arguments passed to fitting function.

Returns:

Dictionary containing: - ‘params’: Fitted parameters - ‘covariance’: Parameter covariance matrix - ‘method’: Method used (‘nlsq’ or ‘scipy’) - ‘success’: Whether fit converged - ‘residuals’: Residuals (y - func(x, *params))

Return type:

dict

Examples

>>> import numpy as np
>>> from piblin_jax.fitting.nlsq import fit_curve
>>>
>>> # Define model function
>>> def linear_model(x, a, b):
...     return a * x + b
>>>
>>> # Generate data
>>> x = np.linspace(0, 10, 50)
>>> y = 2.5 * x + 1.0 + 0.1 * np.random.randn(len(x))
>>>
>>> # Fit
>>> result = fit_curve(linear_model, x, y, p0=[1.0, 0.0])
>>> print(f"Fitted parameters: {result['params']}")
>>> print(f"Method used: {result['method']}")

Notes

NLSQ often provides better convergence than scipy for challenging problems
The scipy fallback ensures the function always works
For Bayesian inference, consider using BayesianModel instead

Curve Fitting

Overview

Quick Examples

Basic Curve Fitting

Automatic Initial Parameter Estimation

Working with Bounds and Constraints

See Also

API Reference

Module Contents

Non-Linear Least Squares