Orchestration Components¶

Added in version 0.6.4.

The orchestration module provides decomposed components for curve fitting, enabling modular testing, customization, and gradual migration via feature flags.

Overview¶

The monolithic CurveFit class has been decomposed into four focused components:

Component	Responsibility
`DataPreprocessor`	Input validation, array conversion, NaN handling
`OptimizationSelector`	Method selection, bounds preparation, initial guess
`CovarianceComputer`	SVD-based covariance, sigma transformation
`StreamingCoordinator`	Memory analysis, streaming strategy selection

DataPreprocessor¶

class nlsq.core.orchestration.DataPreprocessor[source]¶

Bases: object

Preprocessor for curve fitting input data.

Handles: 1. Input validation (type checking, finiteness) 2. Array conversion (numpy/list to JAX) 3. Length consistency checking 4. Data masking for invalid points 5. NaN/Inf handling via nan_policy

Example

>>> preprocessor = DataPreprocessor()
>>> data = preprocessor.preprocess(
...     f=my_model,
...     xdata=x_values,
...     ydata=y_values,
...     sigma=uncertainties,
...     check_finite=True,
... )
>>> print(f"Valid points: {data.n_points}")

preprocess(f, xdata, ydata, *, sigma=None, absolute_sigma=False, check_finite=True, nan_policy='raise', stability_check=False)[source]¶

Validate and preprocess input data for curve fitting.

Parameters:

f (Callable[..., ArrayLike]) – Model function to fit (used for parameter count detection)
xdata (ArrayLike) – Independent variable data
ydata (ArrayLike) – Dependent variable data (observations)
sigma (ArrayLike | None) – Uncertainty/weights for observations
absolute_sigma (bool) – If True, sigma is absolute; else relative
check_finite (bool) – If True, raise on NaN/Inf values
nan_policy (str) – How to handle NaN: ‘raise’, ‘omit’, or ‘propagate’
stability_check (bool) – If True, run additional stability checks

Returns:

PreprocessedData with validated, converted arrays

Raises:

ValueError – If inputs are invalid (wrong shape, non-finite, etc.)
TypeError – If inputs have wrong types

Return type:

PreprocessedData

validate_sigma(sigma, ydata_shape)[source]¶

Validate and convert sigma to appropriate format.

Public interface matching DataPreprocessorProtocol.

Parameters:

sigma (ArrayLike | None) – Input sigma (1D for diagonal, 2D for full covariance)
ydata_shape (tuple[int, ...]) – Shape of ydata for compatibility check

Returns:

Validated numpy array or None

Raises:

ValueError – If sigma shape is incompatible with ydata

Return type:

np.ndarray | None

Example:

from nlsq.core.orchestration import DataPreprocessor

preprocessor = DataPreprocessor()
data = preprocessor.preprocess(
    f=my_model,
    xdata=raw_x,
    ydata=raw_y,
    sigma=uncertainties,
    check_finite=True,
    nan_policy="omit",
)

print(f"Valid points: {data.n_points}")
print(f"NaNs removed: {data.has_nans_removed}")

OptimizationSelector¶

class nlsq.core.orchestration.OptimizationSelector[source]¶

Bases: object

Selector for optimization method and configuration.

Handles: 1. Parameter count detection from function signature 2. Method selection based on bounds and problem type 3. Bounds validation and preparation 4. Initial guess generation if not provided 5. Solver configuration validation

Example

>>> selector = OptimizationSelector()
>>> config = selector.select(
...     f=my_model,
...     xdata=x_values,
...     ydata=y_values,
...     bounds=([0, 0], [10, 10]),
... )
>>> print(f"Method: {config.method}, Params: {config.n_params}")

select(f, xdata, ydata, *, p0=None, bounds=None, method=None, jac=None, tr_solver=None, x_scale=1.0, ftol=1e-08, xtol=1e-08, gtol=1e-08, max_nfev=None)[source]¶

Select optimization method and prepare configuration.

Parameters:

f (Callable[..., ArrayLike]) – Model function to fit
xdata (ArrayLike) – Independent variable data
ydata (ArrayLike) – Dependent variable data
p0 (ArrayLike | None) – Initial parameter guess (auto-detected if None)
bounds (tuple[ArrayLike, ArrayLike] | None) – Parameter bounds as (lower, upper)
method (str | None) – Optimization method (‘trf’, ‘lm’, ‘dogbox’, or None for auto)
jac (str | Callable | None) – Jacobian computation method
tr_solver (str | None) – Trust region solver (‘exact’, ‘lsmr’, or None for auto)
x_scale (ArrayLike | str | float) – Parameter scaling
ftol (float) – Function tolerance
xtol (float) – Parameter tolerance
gtol (float) – Gradient tolerance
max_nfev (int | None) – Maximum function evaluations (auto if None)

Returns:

OptimizationConfig with all settings resolved

Raises:

ValueError – If configuration is invalid

Return type:

OptimizationConfig

detect_parameter_count(f, xdata)[source]¶

Detect number of parameters from function signature.

Uses inspection of function signature to determine parameter count.

Parameters:

f (Callable[..., ArrayLike]) – Model function to analyze
xdata (ArrayLike) – Sample data (not used currently, for future probing)

Returns:

Number of parameters (excluding x)

Raises:

ValueError – If parameter count cannot be determined

Return type:

int

auto_initial_guess(n_params, bounds)[source]¶

Generate automatic initial parameter guess.

Uses bounds midpoint if available, otherwise ones.

Parameters:

n_params (int) – Number of parameters
bounds (tuple[jax.Array, jax.Array] | None) – Parameter bounds or None

Returns:

Initial guess array of shape (n_params,)

Return type:

jax.Array

Example:

from nlsq.core.orchestration import OptimizationSelector

selector = OptimizationSelector()
config = selector.select(
    f=my_model,
    xdata=data.xdata,
    ydata=data.ydata,
    p0=None,  # Auto-detect
    bounds=([0, 0], [10, 10]),
    method=None,  # Auto-select
)

print(f"Method: {config.method}")
print(f"Parameters: {config.n_params}")

CovarianceComputer¶

class nlsq.core.orchestration.CovarianceComputer[source]¶

Bases: object

Computer for parameter covariance from optimization results.

Handles: 1. Jacobian-based covariance via SVD 2. Sigma transformation (1D and 2D) 3. Absolute vs relative sigma handling 4. Singularity detection and handling

Example

>>> computer = CovarianceComputer()
>>> result = computer.compute(
...     result=optimize_result,
...     n_data=100,
...     sigma=uncertainties,
...     absolute_sigma=True,
... )
>>> print(f"Parameter errors: {result.perr}")

__init__()[source]¶

Initialize CovarianceComputer with JIT-compiled functions.

compute(result, n_data, *, sigma=None, absolute_sigma=False, full_output=False)[source]¶

Compute parameter covariance from optimization result.

Uses the Jacobian at the solution to compute covariance via: pcov = (J^T @ J)^(-1) * s_sq

where s_sq is the residual variance.

Parameters:

result (OptimizeResult) – OptimizeResult from LeastSquares
n_data (int) – Number of data points
sigma (jax.Array | None) – Observation uncertainties/weights
absolute_sigma (bool) – If True, sigma is absolute uncertainty
full_output (bool) – If True, include additional diagnostics

Returns:

CovarianceResult with covariance matrix and metadata

Raises:

ValueError – If Jacobian is unavailable or invalid

Return type:

CovarianceResult

create_sigma_transform(sigma, n_data)[source]¶

Create sigma transformation function.

Handles both 1D (diagonal) and 2D (full covariance) sigma.

Parameters:

sigma (jax.Array) – Sigma array, shape (n,) or (n, n)
n_data (int) – Number of data points

Returns:

Tuple of (transform_func, is_2d) - transform_func: Function to apply sigma weighting - is_2d: True if sigma is full covariance matrix

Return type:

tuple[Callable, bool]

compute_condition_number(jacobian)[source]¶

Compute condition number of Jacobian.

Uses singular values: cond = max(s) / min(s)

Parameters:: jacobian (jax.Array) – Jacobian matrix at solution
Returns:: Condition number (inf if singular)
Return type:: float

setup_sigma_transform(sigma, ydata, data_mask, len_diff, m)[source]¶

Setup sigma transformation for weighted least squares.

This is the legacy interface matching CurveFit._setup_sigma_transform.

Parameters:

sigma (np.ndarray | None) – Uncertainty in ydata (1-D errors or 2-D covariance matrix)
ydata (np.ndarray) – Dependent data array
data_mask (np.ndarray) – Boolean mask for valid data points
len_diff (int) – Difference in length for padding
m (int) – Original number of data points

Returns:

Transformation array for sigma or None

Raises:

ValueError – If sigma has incorrect shape or is not positive definite

Return type:

jax.Array | None

Example:

from nlsq.core.orchestration import CovarianceComputer

computer = CovarianceComputer()
cov_result = computer.compute(
    result=optimize_result,
    n_data=len(ydata),
    sigma=uncertainties,
    absolute_sigma=True,
)

print(f"Covariance:\n{cov_result.pcov}")
print(f"Errors: {cov_result.perr}")

StreamingCoordinator¶

class nlsq.core.orchestration.StreamingCoordinator(safety_factor=0.75)[source]¶

Bases: object

Coordinator for streaming strategy selection.

Handles: 1. Memory estimation for dataset + Jacobian 2. Available memory detection 3. Strategy selection based on memory pressure 4. Configuration of chunked/hybrid strategies

Example

>>> coordinator = StreamingCoordinator()
>>> decision = coordinator.decide(
...     xdata=x_array,
...     ydata=y_array,
...     n_params=5,
... )
>>> if decision.strategy == "hybrid":
...     config = decision.hybrid_config
...     # Use hybrid streaming optimizer

__init__(safety_factor=0.75)[source]¶

Initialize StreamingCoordinator.

Parameters:: safety_factor (float) – Memory safety factor (0.75 means use 75% of available)

decide(xdata, ydata, n_params, *, workflow='auto', memory_limit_mb=None, force_streaming=False)[source]¶

Decide on streaming strategy for the dataset.

Analyzes memory requirements and available resources to select the optimal execution strategy.

Parameters:

xdata (jax.Array) – Independent variable data
ydata (jax.Array) – Dependent variable data
n_params (int) – Number of parameters
workflow (str) – Workflow hint (‘auto’, ‘streaming’, ‘hybrid’, ‘normal’)
memory_limit_mb (float | None) – Override for memory limit detection
force_streaming (bool) – If True, always use streaming

Returns:

StreamingDecision with strategy and configuration

Raises:

MemoryError – If dataset too large even for streaming

Return type:

StreamingDecision

estimate_memory(n_data, n_params, dtype_bytes=8)[source]¶

Estimate memory requirement in MB.

Accounts for: - Data arrays (x, y, residuals) - Jacobian matrix (n_data x n_params) - Working arrays for optimization - JAX compilation overhead

Parameters:

n_data (int) – Number of data points
n_params (int) – Number of parameters
dtype_bytes (int) – Bytes per element (8 for float64)

Returns:

Estimated memory in MB

Return type:

float

get_available_memory()[source]¶

Get available system memory in MB.

Cached once per coordinator lifetime (one streaming decision per fit).

Returns:: Available memory in MB
Return type:: float

configure_hybrid(n_data, n_params, available_memory_mb)[source]¶

Configure hybrid streaming for dataset.

Calculates optimal chunk size and strategy parameters.

Parameters:

n_data (int) – Number of data points
n_params (int) – Number of parameters
available_memory_mb (float) – Available memory

Returns:

HybridStreamingConfig for the dataset

Return type:

HybridStreamingConfig

Example:

from nlsq.core.orchestration import StreamingCoordinator

coordinator = StreamingCoordinator()
decision = coordinator.decide(
    xdata=large_x,
    ydata=large_y,
    n_params=5,
    workflow="auto",
)

print(f"Strategy: {decision.strategy}")
print(f"Reason: {decision.reason}")
print(f"Memory pressure: {decision.memory_pressure:.1%}")

Data Classes¶

PreprocessedData¶

class nlsq.interfaces.orchestration_protocol.PreprocessedData(xdata, ydata, sigma, mask, n_points, is_padded, original_length, has_nans_removed, has_infs_removed)[source]¶

Bases: object

Result of data preprocessing.

All arrays are validated and converted to JAX arrays. Padding may be applied for JAX compilation efficiency.

xdata¶

Independent variable data, shape (n,) or (k, n)

Type:: jax.Array

ydata¶

Dependent variable data, shape (n,)

Type:: jax.Array

sigma¶

Uncertainty/weights, shape (n,), (n, n), or None

Type:: jax.Array | None

mask¶

Boolean mask for valid data points, shape (n,)

Type:: jax.Array

n_points¶

Number of valid data points

Type:: int

is_padded¶

Whether arrays were padded for fixed-size compilation

Type:: bool

original_length¶

Length before padding (equals n_points if not padded)

Type:: int

has_nans_removed¶

True if NaN values were filtered during preprocessing

Type:: bool

has_infs_removed¶

True if Inf values were filtered during preprocessing

Type:: bool

xdata: jax.Array¶

ydata: jax.Array¶

sigma: jax.Array | None¶

mask: jax.Array¶

n_points: int¶

is_padded: bool¶

original_length: int¶

has_nans_removed: bool¶

has_infs_removed: bool¶

__init__(xdata, ydata, sigma, mask, n_points, is_padded, original_length, has_nans_removed, has_infs_removed)¶

OptimizationConfig¶

class nlsq.interfaces.orchestration_protocol.OptimizationConfig(method, tr_solver, n_params, p0, bounds, max_nfev, ftol, xtol, gtol, jac, x_scale)[source]¶

Bases: object

Configuration for optimization execution.

Contains all settings needed by LeastSquares optimizer.

method¶

Optimization algorithm (‘trf’, ‘lm’, ‘dogbox’)

Type:: Literal[‘trf’, ‘lm’, ‘dogbox’]

tr_solver¶

Trust region subproblem solver (‘exact’, ‘lsmr’, None)

Type:: Literal[‘exact’, ‘lsmr’] | None

n_params¶

Number of parameters to fit

Type:: int

p0¶

Initial parameter guess

Type:: jax.Array

bounds¶

Lower and upper bounds as (lb, ub) tuple

Type:: tuple[jax.Array, jax.Array]

max_nfev¶

Maximum function evaluations

Type:: int

ftol¶

Relative tolerance for cost function

Type:: float

xtol¶

Relative tolerance for parameters

Type:: float

gtol¶

Relative tolerance for gradient

Type:: float

jac¶

Jacobian specification (‘2-point’, ‘3-point’, callable, None)

Type:: str | Callable | None

x_scale¶

Parameter scaling (‘jac’ or array)

Type:: jax.Array | str

method: Literal['trf', 'lm', 'dogbox']¶

tr_solver: Literal['exact', 'lsmr'] | None¶

n_params: int¶

p0: jax.Array¶

bounds: tuple[jax.Array, jax.Array]¶

max_nfev: int¶

ftol: float¶

xtol: float¶

gtol: float¶

jac: str | Callable | None¶

x_scale: jax.Array | str¶

__init__(method, tr_solver, n_params, p0, bounds, max_nfev, ftol, xtol, gtol, jac, x_scale)¶

CovarianceResult¶

class nlsq.interfaces.orchestration_protocol.CovarianceResult(pcov, perr, method, condition_number, is_singular, sigma_used, absolute_sigma)[source]¶

Bases: object

Result of covariance matrix computation.

pcov¶

Parameter covariance matrix, shape (n, n)

Type:: jax.Array

perr¶

Parameter standard errors (sqrt of diagonal), shape (n,)

Type:: jax.Array

method¶

Computation method used (‘svd’, ‘cholesky’, ‘qr’)

Type:: Literal[‘svd’, ‘cholesky’, ‘qr’]

condition_number¶

Condition number of Jacobian

Type:: float

is_singular¶

True if Jacobian was singular/ill-conditioned

Type:: bool

sigma_used¶

True if sigma weights were applied

Type:: bool

absolute_sigma¶

True if sigma was treated as absolute

Type:: bool

pcov: jax.Array¶

perr: jax.Array¶

method: Literal['svd', 'cholesky', 'qr']¶

condition_number: float¶

is_singular: bool¶

sigma_used: bool¶

absolute_sigma: bool¶

__init__(pcov, perr, method, condition_number, is_singular, sigma_used, absolute_sigma)¶

StreamingDecision¶

class nlsq.interfaces.orchestration_protocol.StreamingDecision(strategy, reason, estimated_memory_mb, available_memory_mb, memory_pressure, chunk_size, n_chunks, hybrid_config)[source]¶

Bases: object

Decision about streaming execution strategy.

strategy¶

Execution strategy to use - ‘direct’: Normal non-streaming fit - ‘chunked’: Simple chunked processing - ‘hybrid’: AdaptiveHybridStreamingOptimizer - ‘auto_memory’: Memory-aware automatic selection

Type:: Literal[‘direct’, ‘chunked’, ‘hybrid’, ‘auto_memory’]

reason¶

Human-readable explanation of decision

Type:: str

estimated_memory_mb¶

Estimated memory requirement

Type:: float

available_memory_mb¶

Available system memory

Type:: float

memory_pressure¶

Memory pressure ratio (0.0 to 1.0)

Type:: float

chunk_size¶

Chunk size for chunked/hybrid strategies

Type:: int | None

n_chunks¶

Number of chunks for chunked strategy

Type:: int | None

hybrid_config¶

Configuration for hybrid strategy

Type:: HybridStreamingConfig | None

strategy: Literal['direct', 'chunked', 'hybrid', 'auto_memory']¶

reason: str¶

estimated_memory_mb: float¶

available_memory_mb: float¶

memory_pressure: float¶

chunk_size: int | None¶

n_chunks: int | None¶

hybrid_config: HybridStreamingConfig | None¶

__init__(strategy, reason, estimated_memory_mb, available_memory_mb, memory_pressure, chunk_size, n_chunks, hybrid_config)¶

Feature Flags¶

Control implementation selection via environment variables for gradual rollout:

# Use old implementation (for rollback)
export NLSQ_PREPROCESSOR_IMPL=old

# Use new implementation (default after rollout)
export NLSQ_PREPROCESSOR_IMPL=new

# Global rollout percentage (0-100)
export NLSQ_REFACTOR_ROLLOUT_PERCENT=50

class nlsq.core.feature_flags.FeatureFlags(preprocessor_impl='auto', selector_impl='auto', covariance_impl='auto', streaming_impl='auto', rollout_percent=0, _session_id=<factory>)[source]¶

Bases: object

Feature flags for CurveFit component extraction.

Controls which implementation (old or new) to use for each extracted component. Supports gradual rollout via percentage-based selection.

preprocessor_impl¶

Implementation choice for DataPreprocessor

Type:: Literal[‘old’, ‘new’, ‘auto’]

selector_impl¶

Implementation choice for OptimizationSelector

Type:: Literal[‘old’, ‘new’, ‘auto’]

covariance_impl¶

Implementation choice for CovarianceComputer

Type:: Literal[‘old’, ‘new’, ‘auto’]

streaming_impl¶

Implementation choice for StreamingCoordinator

Type:: Literal[‘old’, ‘new’, ‘auto’]

rollout_percent¶

Percentage (0-100) of requests using new implementation when impl is ‘auto’

Type:: int

_session_id¶

Internal session ID for deterministic rollout selection

Type:: str

preprocessor_impl: Literal['old', 'new', 'auto']¶

selector_impl: Literal['old', 'new', 'auto']¶

covariance_impl: Literal['old', 'new', 'auto']¶

streaming_impl: Literal['old', 'new', 'auto']¶

rollout_percent: int¶

__post_init__()[source]¶

Validate rollout_percent range.

classmethod from_env(session_id=None)[source]¶

Create FeatureFlags from environment variables.

Parameters:: session_id (str | None) – Optional session ID for deterministic rollout. If not provided, a random ID is generated.
Returns:: FeatureFlags instance with values from environment.
Raises:: ValueError – If environment variable has invalid value.
Return type:: FeatureFlags

get_impl(component)[source]¶

Get implementation choice for a component.

Parameters:: component (str) – Component name (‘preprocessor’, ‘selector’, ‘covariance’, ‘streaming’)
Returns:: Implementation choice for the component
Raises:: ValueError – If component name is unknown
Return type:: Literal[‘old’, ‘new’, ‘auto’]

should_use_new(component)[source]¶

Determine if new implementation should be used.

For ‘old’ or ‘new’ choices, returns the explicit choice. For ‘auto’, uses rollout_percent with deterministic selection based on session_id and component name.

Parameters:: component (str) – Component name (‘preprocessor’, ‘selector’, ‘covariance’, ‘streaming’)
Returns:: True if new implementation should be used, False for old
Raises:: ValueError – If component name is unknown
Return type:: bool

with_override(*, preprocessor_impl=None, selector_impl=None, covariance_impl=None, streaming_impl=None, rollout_percent=None)[source]¶

Create new FeatureFlags with overridden values.

Parameters:

preprocessor_impl (Literal['old', 'new', 'auto'] | None) – Override for preprocessor implementation
selector_impl (Literal['old', 'new', 'auto'] | None) – Override for selector implementation
covariance_impl (Literal['old', 'new', 'auto'] | None) – Override for covariance implementation
streaming_impl (Literal['old', 'new', 'auto'] | None) – Override for streaming implementation
rollout_percent (int | None) – Override for rollout percentage

Returns:

New FeatureFlags instance with overridden values

Return type:

FeatureFlags

to_env_dict()[source]¶

Convert flags to environment variable dictionary.

Returns:: Dictionary mapping env var names to values
Return type:: dict[str, str]

__init__(preprocessor_impl='auto', selector_impl='auto', covariance_impl='auto', streaming_impl='auto', rollout_percent=0, _session_id=<factory>)¶

Orchestration Components¶

Overview¶

DataPreprocessor¶

OptimizationSelector¶

CovarianceComputer¶

StreamingCoordinator¶

Data Classes¶

PreprocessedData¶

OptimizationConfig¶

CovarianceResult¶

StreamingDecision¶

Feature Flags¶

See Also¶