Global Optimization

This reference covers NLSQ’s global optimization capabilities, including multi-start optimization and CMA-ES (Covariance Matrix Adaptation Evolution Strategy).

Overview

NLSQ provides two main approaches for global optimization:

1. Multi-start optimization: Run multiple local optimizations from different starting points using Latin Hypercube Sampling or other quasi-random samplers.

2. CMA-ES (Evolution Strategy): A gradient-free evolutionary algorithm that adapts the search covariance matrix, particularly effective for multi-scale parameter problems.

Installation

Multi-start optimization works out of the box. For CMA-ES, install the optional evosax dependency:

pip install nlsq

CMA-ES Global Optimization

CMA-ES is recommended when:

• Parameters span many orders of magnitude (>1000x scale ratio)

• The fitness landscape has multiple local minima

• Gradient information is unreliable

• You want robust convergence without sensitivity to initialization

Basic Usage

from nlsq.global_optimization import CMAESOptimizer, CMAESConfig
import jax.numpy as jnp


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


# Generate data
x = jnp.linspace(0, 5, 100)
y = 2.5 * jnp.exp(-0.5 * x)

# Bounds are required for CMA-ES
bounds = ([0.1, 0.01], [10.0, 2.0])

# Create optimizer (uses default BIPOP configuration)
optimizer = CMAESOptimizer()

# Run optimization
result = optimizer.fit(model, x, y, bounds=bounds)

print(f"Optimal parameters: {result['popt']}")
print(f"Parameter covariance: {result['pcov']}")

Using Presets

Three presets are available for common use cases:

# Fast preset: no restarts, 50 generations
optimizer = CMAESOptimizer.from_preset("cmaes-fast")

# Standard preset: BIPOP with 9 restarts, 100 generations
optimizer = CMAESOptimizer.from_preset("cmaes")

# Global preset: BIPOP with 9 restarts, 200 generations, 2x population
optimizer = CMAESOptimizer.from_preset("cmaes-global")

Custom Configuration

For fine-grained control, create a custom CMAESConfig:

from nlsq.global_optimization import CMAESConfig, CMAESOptimizer

config = CMAESConfig(
    popsize=32,  # Population size (None = auto)
    max_generations=150,  # Max generations per run
    sigma=0.3,  # Initial step size
    tol_fun=1e-10,  # Fitness tolerance
    tol_x=1e-10,  # Parameter tolerance
    restart_strategy="bipop",  # 'none' or 'bipop'
    max_restarts=5,  # Max BIPOP restarts
    refine_with_nlsq=True,  # Refine with Trust Region
    seed=42,  # For reproducibility
)

optimizer = CMAESOptimizer(config=config)

Memory Management for Large Datasets

CMA-ES can encounter out-of-memory (OOM) issues with large datasets (>10M points) because each fitness evaluation processes the full dataset across all population members. NLSQ provides two strategies to manage memory:

Population Batching: Evaluates population members in smaller groups instead of all at once.

config = CMAESConfig(
    population_batch_size=4,  # Evaluate 4 candidates at a time
)

Data Streaming: Processes the dataset in chunks, accumulating the sum of squared residuals across chunks.

config = CMAESConfig(
    data_chunk_size=50000,  # Process 50K points per chunk
)

Combined Configuration for maximum memory efficiency:

from nlsq.global_optimization import CMAESConfig, CMAESOptimizer

# For 100M+ point datasets
config = CMAESConfig(
    population_batch_size=4,  # Batch population evaluation
    data_chunk_size=50000,  # Stream data in 50K chunks
    max_generations=100,
)

optimizer = CMAESOptimizer(config=config)

Memory Estimation

Use the helper functions to estimate and auto-configure memory usage:

from nlsq.global_optimization import (
    estimate_cmaes_memory_gb,
    auto_configure_cmaes_memory,
)

# Estimate memory for a configuration
memory_gb = estimate_cmaes_memory_gb(
    n_data=100_000_000,
    popsize=16,
    population_batch_size=4,
    data_chunk_size=50000,
)
print(f"Estimated memory: {memory_gb:.3f} GB")  # ~0.005 GB

# Auto-configure based on available memory
pop_batch, data_chunk = auto_configure_cmaes_memory(
    n_data=100_000_000,
    popsize=16,
    available_memory_gb=4.0,  # Target GPU memory
)
config = CMAESConfig(
    population_batch_size=pop_batch,
    data_chunk_size=data_chunk,
)

Memory comparison for 100M points with popsize=16:

Configuration	Peak Memory	Reduction
No batching (default)	12.8 GB	–
population_batch_size=4	3.2 GB	75%
data_chunk_size=50000	~18 MB	99.9%
Both combined	~5 MB	99.96%

Integration with curve_fit

CMA-ES can be used directly through curve_fit with the method parameter:

from nlsq import curve_fit

result = curve_fit(
    model,
    x,
    y,
    bounds=bounds,
    method="cmaes",  # Explicitly request CMA-ES
)

# Or with custom config
from nlsq.global_optimization import CMAESConfig

config = CMAESConfig(max_generations=200, seed=42)
result = curve_fit(
    model,
    x,
    y,
    bounds=bounds,
    method="cmaes",
    cmaes_config=config,
)

Auto Method Selection

Use method="auto" to let NLSQ choose based on the problem:

from nlsq import curve_fit

# NLSQ checks scale ratio and evosax availability
result = curve_fit(model, x, y, bounds=bounds, method="auto")

The MethodSelector class handles the logic:

• If scale ratio > 1000x and evosax available: CMA-ES

• Otherwise: multi-start optimization

Diagnostics

CMA-ES returns detailed diagnostics:

result = optimizer.fit(model, x, y, bounds=bounds)
diag = result["cmaes_diagnostics"]

print(f"Total generations: {diag['total_generations']}")
print(f"Total restarts: {diag['total_restarts']}")
print(f"Final sigma: {diag['final_sigma']}")
print(f"Best fitness: {diag['best_fitness']}")
print(f"Convergence reason: {diag['convergence_reason']}")
print(f"Wall time: {diag['wall_time']}s")

The CMAESDiagnostics class provides analysis methods:

from nlsq.global_optimization import CMAESDiagnostics

diag = CMAESDiagnostics.from_dict(result["cmaes_diagnostics"])
print(diag.summary())
print(f"Fitness improvement: {diag.get_fitness_improvement()}")

BIPOP Restart Strategy

BIPOP (Bi-Population) alternates between large and small population runs:

• Large population: More exploration, broader search

• Small population: More exploitation, faster convergence

The BIPOPRestarter class manages this:

from nlsq.global_optimization import BIPOPRestarter

restarter = BIPOPRestarter(
    base_popsize=16,
    n_params=3,
    max_restarts=9,
    min_fitness_spread=1e-12,
)

while not restarter.exhausted:
    popsize = restarter.get_next_popsize()  # Alternates large/small
    # ... run CMA-ES ...
    restarter.register_restart()

Multi-Start Optimization

For problems where CMA-ES is not needed, multi-start optimization provides robust global search using quasi-random sampling.

See the MultiStartOrchestrator API for details.

API Reference

Configuration

class nlsq.global_optimization.CMAESConfig(popsize=None, max_generations=100, sigma=0.5, tol_fun=1e-08, tol_x=1e-08, restart_strategy='bipop', max_restarts=9, population_batch_size=None, data_chunk_size=None, refine_with_nlsq=True, seed=None)

Bases: object

Configuration for CMA-ES global optimization.

popsize

Population size. If None, uses CMA-ES default: int(4 + 3 * log(n)).

Type:

int | None

max_generations

Maximum number of generations before stopping. Default: 100.

Type:

int

sigma

Initial step size (standard deviation). Default: 0.5.

Type:

float

tol_fun

Function value tolerance for convergence. Default: 1e-8.

Type:

float

tol_x

Parameter tolerance for convergence. Default: 1e-8.

Type:

float

restart_strategy

Restart strategy. Default: ‘bipop’.

Type:

Literal[‘none’, ‘bipop’]

max_restarts

Maximum restart attempts for BIPOP. Default: 9.

Type:

int

population_batch_size

Batch size for population evaluation. If None, evaluates all candidates in parallel. Set to smaller values (e.g., 4) to reduce memory usage.

Type:

int | None

data_chunk_size

Chunk size for data streaming. If None, processes full dataset at once. Set to smaller values (e.g., 50000) for datasets >10M points to avoid OOM. Must be >= 1024 for numerical stability.

Type:

int | None

refine_with_nlsq

Whether to refine best solution with NLSQ TRF. Default: True.

Type:

bool

seed

Random seed for reproducibility. If None, uses random seed.

Type:

int | None

Examples

>>> config = CMAESConfig(popsize=32, max_generations=200)
>>> config = CMAESConfig.from_preset('cmaes-global')

Memory-efficient configuration for large datasets:

>>> config = CMAESConfig(
...     population_batch_size=4,  # Evaluate 4 candidates at a time
...     data_chunk_size=50000,    # Process 50K points per chunk
... )

popsize: int | None

max_generations: int

sigma: float

tol_fun: float

tol_x: float

restart_strategy: Literal['none', 'bipop']

max_restarts: int

population_batch_size: int | None

data_chunk_size: int | None

refine_with_nlsq: bool

seed: int | None

__post_init__()

Validate configuration after initialization.

classmethod from_preset(preset_name)

Create a CMAESConfig from a named preset.

Parameters:

preset_name (str) – Name of the preset. One of ‘cmaes-fast’, ‘cmaes’, ‘cmaes-global’.

Returns:

Configuration for the specified preset.

Return type:

CMAESConfig

Raises:

ValueError – If preset_name is not recognized.

Examples

>>> config = CMAESConfig.from_preset('cmaes-fast')
>>> config.max_generations
50

__init__(popsize=None, max_generations=100, sigma=0.5, tol_fun=1e-08, tol_x=1e-08, restart_strategy='bipop', max_restarts=9, population_batch_size=None, data_chunk_size=None, refine_with_nlsq=True, seed=None)

Optimizer

class nlsq.global_optimization.CMAESOptimizer(config=None)

Bases: object

CMA-ES global optimizer with NLSQ refinement using evosax.

Uses evosax’s CMA-ES implementation for gradient-free global optimization, followed by NLSQ Trust Region Reflective refinement for proper parameter covariance estimation.

Parameters:

config (CMAESConfig | None, optional) – Configuration for CMA-ES optimization. If None, uses default config.

config

Configuration for CMA-ES optimization.

Type:

CMAESConfig

Examples

>>> from nlsq.global_optimization import CMAESOptimizer, CMAESConfig
>>> import jax.numpy as jnp
>>>
>>> def model(x, a, b):
...     return a * jnp.exp(-b * x)
>>>
>>> x = jnp.linspace(0, 5, 100)
>>> y = 2.5 * jnp.exp(-0.5 * x)
>>> bounds = ([0.1, 0.01], [10.0, 2.0])
>>>
>>> optimizer = CMAESOptimizer()
>>> result = optimizer.fit(model, x, y, bounds=bounds)
>>> print(f"Optimal params: {result['popt']}")

__init__(config=None)

Initialize CMAESOptimizer.

Parameters:

config (CMAESConfig | None, optional) – Configuration for CMA-ES optimization. If None, uses default config (BIPOP enabled, 100 generations, 9 max restarts).

classmethod from_preset(preset_name)

Create optimizer from a named preset.

Parameters:

preset_name (str) – Name of the preset. One of ‘cmaes-fast’, ‘cmaes’, ‘cmaes-global’.

Returns:

Optimizer configured with the specified preset.

Return type:

CMAESOptimizer

Examples

>>> optimizer = CMAESOptimizer.from_preset('cmaes-fast')
>>> optimizer.config.max_generations
50

fit(f, xdata, ydata, p0=None, bounds=None, sigma=None, **kwargs)

Run CMA-ES global optimization followed by NLSQ refinement.

Parameters:

• f (Callable) – Model function f(x, *params) -> y.

• xdata (ArrayLike) – Independent variable data.

• ydata (ArrayLike) – Dependent variable data.

• p0 (ArrayLike | None, optional) – Initial parameter guess. If None, uses center of bounds.

• bounds (tuple[ArrayLike, ArrayLike] | None) – Lower and upper bounds for parameters. Required for CMA-ES.

• sigma (ArrayLike | None, optional) – Standard deviation of ydata for weighted residuals.

• **kwargs (Any) – Additional keyword arguments (passed to NLSQ refinement).

Returns:

Result dictionary containing: - popt: Optimal parameters - pcov: Parameter covariance matrix (from NLSQ refinement) - Additional fields from NLSQ result

Return type:

dict[str, Any]

Raises:

ValueError – If bounds are not provided (required for CMA-ES).

Diagnostics

class nlsq.global_optimization.CMAESDiagnostics(total_generations=0, total_restarts=0, final_sigma=0.0, best_fitness=inf, fitness_history=<factory>, restart_history=<factory>, convergence_reason='not_converged', nlsq_refinement=False, wall_time=0.0)

Bases: object

Diagnostic information from CMA-ES optimization.

total_generations

Total number of CMA-ES generations across all restarts.

Type:

int

total_restarts

Number of BIPOP restarts performed.

Type:

int

final_sigma

Final step size (standard deviation) at convergence.

Type:

float

best_fitness

Best fitness value found (negative SSR, higher is better).

Type:

float

fitness_history

History of best fitness values per generation.

Type:

list[float]

restart_history

Information about each restart (popsize, generations, reason).

Type:

list[dict[str, Any]]

convergence_reason

Reason for convergence or termination.

Type:

str

nlsq_refinement

Whether NLSQ refinement was applied.

Type:

bool

wall_time

Total wall-clock time in seconds.

Type:

float

Examples

>>> diagnostics = CMAESDiagnostics(
...     total_generations=150,
...     total_restarts=3,
...     final_sigma=0.01,
...     best_fitness=-1e-10,
...     convergence_reason="fitness_tolerance",
... )
>>> print(diagnostics.summary())

total_generations: int

total_restarts: int

final_sigma: float

best_fitness: float

fitness_history: list[float]

restart_history: list[dict[str, Any]]

convergence_reason: str

nlsq_refinement: bool

wall_time: float

summary()

Generate a human-readable summary of the diagnostics.

Returns:

Multi-line summary string.

Return type:

str

to_dict()

Convert diagnostics to a dictionary.

Returns:

Dictionary representation for serialization.

Return type:

dict

classmethod from_dict(d)

Create diagnostics from a dictionary.

Parameters:

d (dict) – Dictionary with diagnostics values.

Returns:

Diagnostics instance.

Return type:

CMAESDiagnostics

get_fitness_improvement()

Calculate fitness improvement from first to last generation.

Returns:

Relative fitness improvement, or None if not enough history.

Return type:

float | None

get_convergence_rate()

Compute per-generation convergence rate.

Returns:

Array of per-generation fitness changes, or None if not enough history.

Return type:

NDArray | None

__init__(total_generations=0, total_restarts=0, final_sigma=0.0, best_fitness=inf, fitness_history=<factory>, restart_history=<factory>, convergence_reason='not_converged', nlsq_refinement=False, wall_time=0.0)

Restart Strategy

class nlsq.global_optimization.BIPOPRestarter(base_popsize, n_params, max_restarts=9, min_fitness_spread=1e-12)

Bases: object

BIPOP restart manager for CMA-ES optimization.

Manages alternating large/small population restarts following the BIPOP strategy. Large populations explore broadly while small populations exploit local regions more intensively.

Parameters:

• base_popsize (int) – Base population size (typically 4 + floor(3 * ln(n))).

• n_params (int) – Number of parameters being optimized.

• max_restarts (int, optional) – Maximum number of restarts before giving up. Default is 9.

• min_fitness_spread (float, optional) – Minimum fitness spread threshold for stagnation detection. Default is 1e-12.

restart_count

Number of restarts performed so far.

Type:

int

exhausted

True if max_restarts has been reached.

Type:

bool

best_solution

Best solution found across all restarts.

Type:

jax.Array | None

best_fitness

Best fitness found across all restarts.

Type:

float

Examples

>>> restarter = BIPOPRestarter(base_popsize=8, n_params=3)
>>> popsize = restarter.get_next_popsize()
>>> # ... run CMA-ES with popsize ...
>>> if restarter.check_stagnation(fitness_spread=1e-15):
...     restarter.register_restart()
...     popsize = restarter.get_next_popsize()

base_popsize: int

n_params: int

max_restarts: int = 9

min_fitness_spread: float = 1e-12

restart_count: int = 0

property exhausted: bool

Whether maximum restarts have been reached.

property best_solution: jax.Array | None

Best solution found across all restarts.

property best_fitness: float

Best fitness found across all restarts.

get_next_popsize()

Get population size for the next run.

Returns:

Population size to use for next CMA-ES run. Large runs use 2x base_popsize, small runs use base_popsize/2 to base_popsize.

Return type:

int

register_restart()

Register that a restart has occurred.

Call this after a restart to update internal state.

check_stagnation(fitness_spread)

Check if optimization has stagnated.

Parameters:

fitness_spread (float) – Difference between max and min fitness in current population.

Returns:

True if stagnation detected (fitness spread below threshold).

Return type:

bool

update_best(solution, fitness)

Update best solution if the new one is better.

Parameters:

• solution (jax.Array) – Candidate solution.

• fitness (float) – Fitness value (CMA-ES maximizes, so higher is better).

get_best()

Get the best solution found across all restarts.

Returns:

Best solution and its fitness, or (None, -inf) if none found.

Return type:

tuple[jax.Array | None, float]

reset()

Reset the restarter to initial state.

__init__(base_popsize, n_params, max_restarts=9, min_fitness_spread=1e-12)

Method Selection

class nlsq.global_optimization.MethodSelector(scale_threshold=1000.0)

Bases: object

Select optimization method based on problem characteristics.

The selector analyzes the parameter bounds to compute a scale ratio, which indicates how many orders of magnitude separate the parameter scales. CMA-ES is preferred for multi-scale problems (high scale ratio) when evosax is available.

Parameters:

scale_threshold (float, optional) – Threshold for scale ratio above which CMA-ES is preferred. Default is 1000.0 (3 orders of magnitude).

scale_threshold

The scale ratio threshold for CMA-ES preference.

Type:

float

Examples

>>> from nlsq.global_optimization import MethodSelector
>>> import numpy as np
>>>
>>> selector = MethodSelector()
>>> lower = np.array([1e2, 1e-5, 0.5])
>>> upper = np.array([1e6, 1e-1, 3.0])
>>>
>>> method = selector.select('auto', lower, upper)
>>> print(f"Selected method: {method}")

__init__(scale_threshold=1000.0)

Initialize MethodSelector.

Parameters:

scale_threshold (float, optional) – Threshold for scale ratio above which CMA-ES is preferred.

compute_scale_ratio(lower_bounds, upper_bounds)

Compute the scale ratio from parameter bounds.

The scale ratio is the ratio of the largest parameter range to the smallest, indicating how many orders of magnitude separate the parameter scales.

Parameters:

• lower_bounds (ArrayLike) – Lower bounds for parameters.

• upper_bounds (ArrayLike) – Upper bounds for parameters.

Returns:

Scale ratio (>= 1.0). Higher values indicate more diverse scales.

Return type:

float

select(requested_method, lower_bounds, upper_bounds)

Select the optimization method to use.

Parameters:

• requested_method ({'cmaes', 'multi-start', 'auto'} | None) – Requested method. If ‘auto’ or None, selection is based on scale ratio and evosax availability.

• lower_bounds (ArrayLike) – Lower bounds for parameters.

• upper_bounds (ArrayLike) – Upper bounds for parameters.

Returns:

The selected optimization method.

Return type:

Literal[‘cmaes’, ‘multi-start’]

Notes

Selection logic:

1. If ‘cmaes’ requested and evosax available -> ‘cmaes’

2. If ‘cmaes’ requested but evosax unavailable -> ‘multi-start’ (with warning)

3. If ‘multi-start’ requested -> ‘multi-start’

4. If ‘auto’ or None:

   • If scale_ratio > threshold and evosax available -> ‘cmaes’

   • Otherwise -> ‘multi-start’

Memory Helpers

nlsq.global_optimization.estimate_cmaes_memory_gb(n_data, popsize, population_batch_size=None, data_chunk_size=None)

Estimate peak GPU memory usage for CMA-ES in GB.

Parameters:

• n_data (int) – Number of data points.

• popsize (int) – Population size.

• population_batch_size (int | None, optional) – Batch size for population evaluation. If None, uses full popsize.

• data_chunk_size (int | None, optional) – Chunk size for data streaming. If None, uses full dataset.

Returns:

Estimated peak memory usage in GB.

Return type:

float

Examples

>>> estimate_cmaes_memory_gb(10_000_000, popsize=16)
1.1920928955078125

>>> estimate_cmaes_memory_gb(10_000_000, popsize=16, population_batch_size=4)
0.298023223876953

>>> estimate_cmaes_memory_gb(10_000_000, popsize=16, data_chunk_size=50000)
0.005960464477539062

nlsq.global_optimization.auto_configure_cmaes_memory(n_data, popsize, available_memory_gb=8.0, safety_factor=0.7)

Auto-configure batch sizes to fit in available memory.

Parameters:

• n_data (int) – Number of data points.

• popsize (int) – Population size.

• available_memory_gb (float, optional) – Available GPU/CPU memory in GB. Default: 8.0.

• safety_factor (float, optional) – Safety factor for memory allocation (0-1). Default: 0.7.

Returns:

(population_batch_size, data_chunk_size) configuration. None means no batching/chunking needed for that dimension.

Return type:

tuple[int | None, int | None]

Examples

>>> auto_configure_cmaes_memory(1_000_000, popsize=16, available_memory_gb=8.0)
(None, None)  # No batching needed

>>> auto_configure_cmaes_memory(100_000_000, popsize=16, available_memory_gb=8.0)
(4, None)  # Population batching only

>>> auto_configure_cmaes_memory(100_000_000, popsize=16, available_memory_gb=1.0)
(1, 65536)  # Both population batching and data chunking