nlsq.callbacks module

Progress callbacks for monitoring optimization iterations.

This module provides built-in callbacks for monitoring curve_fit optimization progress, including progress bars, logging, early stopping, and live plotting.

Callbacks are called after each optimization iteration with information about the current state: iteration number, cost, parameters, gradient norm, etc.

class nlsq.callbacks.CallbackBase[source]

Bases: object

Base class for optimization callbacks.

Subclass this to create custom callbacks. Override the __call__ method to define what happens at each iteration.

Examples

>>> class CustomCallback(CallbackBase):
...     def __call__(self, iteration, cost, params, info):
...         print(f"Iter {iteration}: cost={cost:.6f}")
__call__(iteration, cost, params, info)[source]

Called after each optimization iteration.

Parameters:
  • iteration (int) – Current iteration number (0-indexed)

  • cost (float) – Current cost/objective function value

  • params (np.ndarray) – Current parameter values

  • info (dict) – Additional information (gradient_norm, nfev, etc.)

close()[source]

Clean up resources.

Override this method if your callback uses resources that need explicit cleanup (files, network connections, etc.).

class nlsq.callbacks.CallbackChain(*callbacks)[source]

Bases: CallbackBase

Chain multiple callbacks together.

Calls each callback in sequence. If any callback raises StopOptimization, propagates it to stop the optimization.

Parameters:

*callbacks (CallbackBase) – Callbacks to chain together

Examples

>>> from nlsq.callbacks import CallbackChain, ProgressBar, EarlyStopping
>>> callback = CallbackChain(
...     ProgressBar(max_nfev=100),
...     EarlyStopping(patience=5)
... )
>>> popt, pcov = curve_fit(f, x, y, callback=callback)
__init__(*callbacks)[source]
__call__(iteration, cost, params, info)[source]

Call all callbacks in sequence.

close()[source]

Close all callbacks that have a close method.

__del__()[source]

Ensure all callbacks are closed.

class nlsq.callbacks.EarlyStopping(patience=10, min_delta=1e-06, verbose=True)[source]

Bases: CallbackBase

Stop optimization early if no improvement for patience iterations.

Parameters:
  • patience (int, optional) – Number of iterations with no improvement to wait before stopping. Default: 10

  • min_delta (float, optional) – Minimum change in cost to qualify as an improvement. Default: 1e-6

  • verbose (bool, optional) – Whether to print messages. Default: True

Examples

>>> from nlsq.callbacks import EarlyStopping
>>> callback = EarlyStopping(patience=5, min_delta=1e-4)
>>> popt, pcov = curve_fit(f, x, y, callback=callback)

Notes

Raises StopOptimization exception when patience is exceeded, which will be caught by the optimizer and treated as successful convergence.

__init__(patience=10, min_delta=1e-06, verbose=True)[source]
__call__(iteration, cost, params, info)[source]

Check for improvement and stop if patience exceeded.

class nlsq.callbacks.IterationLogger(filename=None, mode='w', log_params=False, file=None)[source]

Bases: CallbackBase

Log optimization progress to file or stdout.

Parameters:
  • filename (str, optional) – File to log to. If None and file is None, logs to stdout.

  • mode (str, optional) – File open mode. Default: ‘w’ (overwrite)

  • log_params (bool, optional) – Whether to log parameter values. Default: False

  • file (file-like object, optional) – File-like object to write to. If provided, filename is ignored.

Examples

>>> from nlsq.callbacks import IterationLogger
>>> callback = IterationLogger("optimization.log")
>>> popt, pcov = curve_fit(f, x, y, callback=callback)
__init__(filename=None, mode='w', log_params=False, file=None)[source]
__call__(iteration, cost, params, info)[source]

Log iteration information.

close()[source]

Close log file.

__del__()[source]

Ensure file is closed on deletion.

class nlsq.callbacks.ProgressBar(max_nfev=None, desc='Optimizing', **tqdm_kwargs)[source]

Bases: CallbackBase

Progress bar callback using tqdm.

Displays a progress bar showing optimization progress with current cost, gradient norm, and iteration statistics.

Parameters:
  • max_nfev (int, optional) – Maximum number of function evaluations. If provided, progress bar will be based on nfev. Otherwise, shows indefinite progress.

  • desc (str, optional) – Description to display in progress bar. Default: “Optimizing”

  • **tqdm_kwargs – Additional keyword arguments passed to tqdm

Examples

>>> from nlsq import curve_fit
>>> from nlsq.callbacks import ProgressBar
>>> callback = ProgressBar(max_nfev=100)
>>> popt, pcov = curve_fit(f, x, y, callback=callback)
__init__(max_nfev=None, desc='Optimizing', **tqdm_kwargs)[source]
__call__(iteration, cost, params, info)[source]

Update progress bar.

close()[source]

Close progress bar.

__del__()[source]

Ensure progress bar is closed on deletion.

exception nlsq.callbacks.StopOptimization[source]

Bases: Exception

Exception raised by callbacks to request early termination.

Overview

The nlsq.callbacks module provides callback functions for monitoring and controlling optimization progress. Callbacks can display progress bars, log iterations, implement early stopping, and save intermediate results.

New in version 0.1.1: Complete callback system with progress bars and early stopping.

Key Features

  • Real-time progress monitoring with rich progress bars

  • Early stopping based on convergence criteria

  • Iteration logging for debugging and analysis

  • Result saving for checkpointing

  • Custom callbacks for specialized monitoring

Classes

CallbackBase()

Base class for optimization callbacks.

CallbackChain(*callbacks)

Chain multiple callbacks together.

StopOptimization

Exception raised by callbacks to request early termination.

ProgressBar([max_nfev, desc])

Progress bar callback using tqdm.

EarlyStopping([patience, min_delta, verbose])

Stop optimization early if no improvement for patience iterations.

IterationLogger([filename, mode, ...])

Log optimization progress to file or stdout.

Usage Examples

Progress Bar

Display a real-time progress bar during optimization:

from nlsq import curve_fit
from nlsq.callbacks import ProgressBar
import jax.numpy as jnp


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


result = curve_fit(model, x, y, p0=[1.0, 0.5], callback=ProgressBar())

Early Stopping

Stop optimization when parameters converge:

from nlsq.callbacks import EarlyStopping

callback = EarlyStopping(patience=10, min_delta=1e-6, monitor="cost")

result = curve_fit(model, x, y, p0=[1.0, 0.5], callback=callback)

print(f"Stopped at iteration {callback.stopped_iteration}")

Iteration Logging

Log detailed information at each iteration:

from nlsq.callbacks import IterationLogger

logger = IterationLogger(log_every=5, log_cost=True, log_gradient_norm=True)

result = curve_fit(model, x, y, p0=[1.0, 0.5], callback=logger)

# Access logged history
print(logger.history)

Result Saving

Save intermediate results during optimization:

from nlsq.callbacks import ResultSaver

saver = ResultSaver(save_every=10, save_path="checkpoints/")

result = curve_fit(model, x, y, p0=[1.0, 0.5], callback=saver)

Combining Callbacks

Use multiple callbacks simultaneously:

from nlsq.callbacks import CallbackComposer, ProgressBar, EarlyStopping

callbacks = CallbackComposer(
    [ProgressBar(), EarlyStopping(patience=10), IterationLogger(log_every=5)]
)

result = curve_fit(model, x, y, p0=[1.0, 0.5], callback=callbacks)

Custom Callbacks

Create your own custom callback by subclassing OptimizationCallback:

from nlsq.callbacks import OptimizationCallback


class CustomCallback(OptimizationCallback):
    def on_iteration_end(self, iteration, params, cost, gradient_norm):
        if cost < 1e-10:
            print("Excellent fit achieved!")
            self.stop_training = True


result = curve_fit(model, x, y, p0=[1.0, 0.5], callback=CustomCallback())

Callback Methods

All callbacks inherit from OptimizationCallback and can override these methods:

  • on_optimization_begin(): Called before optimization starts

  • on_iteration_begin(iteration): Called before each iteration

  • on_iteration_end(iteration, params, cost, gradient_norm): Called after each iteration

  • on_optimization_end(result): Called after optimization completes

Interactive Notebooks

  • Callbacks Demo (15 min) - Progress monitoring, early stopping, and custom callbacks

See Also