Optimization Module

The homodyne.optimization module provides two complementary optimization methods for parameter estimation in homodyne scattering analysis:

1. NLSQ (Primary): Fast, reliable trust-region optimization using Levenberg-Marquardt

2. CMC (Secondary): Bayesian uncertainty quantification using Consensus Monte Carlo

Overview

Optimization Philosophy:

• NLSQ as primary method for fast parameter estimation

• CMC (NumPyro/NUTS) for publication-quality uncertainty quantification

• CPU-optimized architecture

• Dataset size-aware strategy selection

Method Comparison:

Method	Use Case	Performance
NLSQ	Fast parameter estimation, exploratory analysis	~seconds for 1M points
CMC	Uncertainty quantification, publication figures	~minutes (parallelized)

Module Contents

JAX-First Optimization for Homodyne.4

Simplified optimization system using NLSQ package (primary) and CMC (high-accuracy Bayesian) for robust parameter estimation in homodyne analysis.

This module implements the streamlined optimization philosophy: 1. NLSQ as primary method (fast, reliable parameter estimation) 2. CMC (NumPyro/NUTS) for uncertainty quantification 3. Unified homodyne model: c2_fitted = c2_theory * contrast + offset

Key Features: - NLSQ trust-region optimization (Levenberg-Marquardt) as foundation - CMC: Fresh reimplementation with ArviZ-native output - CPU-primary architecture (GPU removed in v2.3.0) - Dataset size-aware optimization strategies

Performance Comparison: - NLSQ: Fast, reliable parameter estimation - CMC: Full posterior sampling, publication-quality uncertainty

Note: Legacy mcmc/ package removed in v3.0. CMC is the sole MCMC backend.

homodyne.optimization.get_optimization_info()

Get information about available optimization methods.

Returns:

Dictionary with availability status and recommendations

Return type:

dict[str, Any]

Primary Functions

homodyne.optimization.fit_nlsq_jax	NLSQ trust-region nonlinear least squares optimization with per-angle scaling.
homodyne.optimization.fit_mcmc_jax	Run CMC (Consensus Monte Carlo) analysis on XPCS data.
homodyne.optimization.get_optimization_info	Get information about available optimization methods.

NLSQ: Non-Linear Least Squares

Trust-region optimization using the Levenberg-Marquardt algorithm, implemented via the nlsq package.

Core Module

NLSQ: Primary Optimization Method for Homodyne

NLSQ package-based trust-region nonlinear least squares solver for the scaled optimization process. This is the primary optimization method providing fast, reliable parameter estimation for homodyne analysis.

Core Equation: c₂(φ,t₁,t₂) = 1 + contrast × [c₁(φ,t₁,t₂)]²

Key Features: - NLSQ trust-region solver (TRF/Levenberg-Marquardt) for robust optimization - JAX JIT compilation for high performance - Intelligent error recovery with 3-attempt retry strategy (T022-T024) - Compatible with existing ParameterSpace and FitResult classes - HPC-optimized for 36/128-core CPU nodes - CPU-only (no GPU support since v2.3.0) - Dataset size-aware optimization strategies

Performance (Validated T036-T041): - Parameter recovery accuracy: 2-14% on core parameters - Sub-linear time scaling: ~1.5s for 500-9,375 point datasets - Numerical stability: <4% deviation across initial conditions - Throughput: 317-5,977 points/second - 100% convergence rate across all validation tests

Production Status: - Scientifically validated (7/7 tests passed) - Production-ready with error recovery - Approved for scientific research and deployment

Migration from Optimistix: - Replaced Optimistix with NLSQ package (github.com/imewei/NLSQ) - NLSQWrapper provides unified interface with error recovery - Maintains backward API compatibility - User-facing Optimistix references removed from public APIs

References: - NLSQ Package: https://github.com/imewei/NLSQ - Validation Report: SCIENTIFIC_VALIDATION_REPORT.md - Production Report: PRODUCTION_READINESS_REPORT.md

class homodyne.optimization.nlsq.core.NLSQResult

Bases: object

Result container for NLSQ optimization compatible with FitResult.

__init__(parameters, parameter_errors, chi_squared, reduced_chi_squared, success, message, n_iterations, optimization_time, method='nlsq')

homodyne.optimization.nlsq.core.fit_nlsq_jax(data, config, initial_params=None, per_angle_scaling=True, use_adapter=False, _skip_global_selection=False)

NLSQ trust-region nonlinear least squares optimization with per-angle scaling.

Uses NLSQ package (github.com/imewei/NLSQ) for trust-region optimization.

v2.11.0+: Experimental NLSQAdapter with CurveFit class available for improved JIT caching and automatic workflow selection. Set use_adapter=True to enable (default is False, uses NLSQWrapper).

Primary optimization method implementing the scaled optimization process: c₂(φ,t₁,t₂) = 1 + contrast × [c₁(φ,t₁,t₂)]²

Parameters:

• data (dict[str, Any]) –

  XPCS experimental data. Accepts two formats:

  Format 1 (CLI/loader format): - ‘phi_angles_list’: phi angle array (mapped to ‘phi’) - ‘c2_exp’: experimental correlation data (n_phi, n_t1, n_t2) (mapped to ‘g2’) - ‘t1’: first delay time array - ‘t2’: second delay time array - ‘wavevector_q_list’: q-vector array (first element extracted as scalar ‘q’) - ‘sigma’: (optional) uncertainty array, defaults to 0.01 * ones_like(g2) - ‘L’: (optional) stator-rotor gap (rheology) or sample-detector distance (standard XPCS), defaults to config value or 2000000 Å (200 µm, typical rheology-XPCS gap) - ‘dt’: (optional) time step, defaults to config value or None

  Format 2 (Direct format): - ‘phi’: phi angle array - ‘g2’: experimental correlation data (n_phi, n_t1, n_t2) - ‘t1’: first delay time array - ‘t2’: second delay time array - ‘q’: wavevector magnitude (scalar) - ‘sigma’: (optional) uncertainty array - ‘L’: (optional) stator-rotor gap or sample-detector distance [Å] - ‘dt’: (optional) time step [s]

• config (ConfigManager) – Configuration manager with optimization settings

• initial_params (dict[str, float] | None) – Initial parameter guesses. If None, uses defaults from config.

• per_angle_scaling (bool) – MUST be True. Per-angle contrast/offset parameters are physically correct as each scattering angle has different optical properties and detector responses. Legacy scalar mode (False) is no longer supported (removed Nov 2025).

• use_adapter (bool) – EXPERIMENTAL (v2.11.0+): If True, use NLSQAdapter with NLSQ’s CurveFit class for improved JIT caching and automatic workflow selection. If False (default), use the stable NLSQWrapper implementation.

Notes

Global Optimization Selection (v2.15.0+): This function serves as the unified entry point for NLSQ optimization. When called, it first checks for global optimization methods:

1. If cmaes.enable: true → delegates to fit_nlsq_cmaes()

2. If multi_start.enable: true → delegates to fit_nlsq_multistart()

3. Otherwise → runs local trust-region optimization

The CMA-ES function will fall back to multi-start (if enabled) when the scale ratio is below the threshold, implementing the full fallback chain.

Returns:

Optimization result with parameters, uncertainties, and diagnostics

Return type:

OptimizationResult

Raises:

• ImportError – If NLSQ package is not available

• ValueError – If data validation fails

homodyne.optimization.nlsq.core.fit_nlsq_multistart(data, config, initial_params=None, per_angle_scaling=True)

Multi-start NLSQ optimization with Latin Hypercube Sampling.

This function explores the parameter space using Latin Hypercube Sampling to avoid local minima. FULL strategy is always used regardless of dataset size - numerical precision and reproducibility take priority over speed.

NOTE: Subsampling is explicitly NOT supported per project requirements.

Parameters:

• data (dict[str, Any]) – XPCS experimental data with keys: - wavevector_q_list: Q-vector values - phi_angles_list: Azimuthal angles - t1, t2: Time coordinates - c2_exp: Experimental g2 correlation data - sigma (optional): Error weights

• config (ConfigManager) – Configuration manager with optimization.nlsq.multi_start settings.

• initial_params (dict[str, float] | None) – Initial parameter guess. If provided, included as one of the starts.

• per_angle_scaling (bool) – Whether to use per-angle contrast/offset scaling. Default: True.

Returns:

Aggregated results including: - best: Best result by chi-squared - all_results: All optimization attempts - strategy_used: “full” (only supported strategy) - n_unique_basins: Number of distinct local minima found - degeneracy_detected: Whether parameter degeneracy was detected

Return type:

MultiStartResult

Raises:

• ImportError – If multi-start module is not available.

• ValueError – If multi-start is not enabled in configuration.

Examples

>>> config = ConfigManager("config.yaml")
>>> # Ensure multi_start.enable: true in config
>>> result = fit_nlsq_multistart(data, config)
>>> print(f"Best chi2: {result.best.chi_squared:.4g}")
>>> print(f"Strategy used: {result.strategy_used}")
>>> if result.degeneracy_detected:
...     print(f"Warning: {result.n_unique_basins} distinct basins found")

homodyne.optimization.nlsq.core.fit_nlsq_cmaes(data, config, initial_params=None, per_angle_scaling=True)

CMA-ES global optimization for multi-scale parameter problems.

Uses NLSQ’s CMAESOptimizer with evosax backend for global optimization. Particularly beneficial for laminar_flow mode where parameters have vastly different scales (e.g., D₀ ~ 1e4 vs γ̇₀ ~ 1e-3, scale ratio > 1e7).

Features: - Covariance Matrix Adaptation for multi-scale parameters - BIPOP restart strategy for robust convergence - Memory batching/streaming for large datasets - Optional L-M refinement of CMA-ES solution

Parameters:

• data (dict[str, Any]) – XPCS experimental data (same format as fit_nlsq_jax).

• config (ConfigManager) – Configuration manager with optimization.nlsq.cmaes settings.

• initial_params (dict[str, float] | None) – Initial parameter guess. Used as CMA-ES starting point.

• per_angle_scaling (bool) – Whether to use per-angle contrast/offset scaling. Default: True.

Returns:

Optimization result with parameters, uncertainties, and diagnostics.

Return type:

OptimizationResult

Raises:

• ImportError – If CMA-ES is not available (requires NLSQ 0.6.4+ with evosax).

• ValueError – If CMA-ES is not enabled in configuration.

Examples

>>> config = ConfigManager("config.yaml")
>>> # Ensure cmaes.enable: true in config
>>> result = fit_nlsq_cmaes(data, config)
>>> print(f"Chi2: {result.chi_squared:.4e}")
>>> print(f"Method: {result.device_info['method']}")

Wrapper (Legacy)

High-level interface with automatic strategy selection.

NLSQ Wrapper for Homodyne Optimization.

Role and When to Use (v2.11.0+)

NLSQWrapper (this module) is the stable fallback adapter for: - Complex optimizations requiring full anti-degeneracy integration - laminar_flow mode with many phi angles (> 6) - Large datasets (> 100M points) requiring streaming/chunking strategies - Custom transforms or advanced recovery mechanisms - Production stability where reliability is critical

Use NLSQAdapter instead for: - Standard optimizations (static_isotropic mode) - Small to medium datasets (< 10M points) - Multi-start optimization (model caching provides 3-5× speedup) - Performance-critical workflows requiring JIT compilation

Key Differences:

• Model caching: NLSQWrapper=None, NLSQAdapter=Built-in

• JIT compilation: NLSQWrapper=Manual, NLSQAdapter=Auto

• Workflow auto-select: NLSQWrapper=Custom, NLSQAdapter=Via NLSQ

• Anti-degeneracy layers: NLSQWrapper=Full, NLSQAdapter=Via fit()

• Recovery system: NLSQWrapper=3-attempt, NLSQAdapter=NLSQ native

• Streaming support: NLSQWrapper=Full custom, NLSQAdapter=Via NLSQ

Decision Guide:

1. If you need robust streaming for 100M+ points: Use NLSQWrapper

2. If you need full anti-degeneracy control: Use NLSQWrapper

3. If you need maximum speed for multi-start optimization: Use NLSQAdapter

4. Default recommendation: NLSQAdapter with automatic fallback to NLSQWrapper

This module provides an adapter layer between homodyne’s optimization API and the NLSQ package’s trust-region nonlinear least squares interface.

The NLSQWrapper class implements the Adapter pattern to translate: - Homodyne’s multi-dimensional XPCS data → NLSQ’s flattened array format - Homodyne’s parameter bounds tuple → NLSQ’s (lower, upper) format - NLSQ’s (popt, pcov) output → Homodyne’s OptimizationResult dataclass

Key Features: - Automatic dataset size detection and strategy selection - Angle-stratified chunking for per-angle parameter compatibility (v2.2+) - Intelligent error recovery with 3-attempt retry strategy (T022-T024) - Actionable error diagnostics with 5 error categories - CPU-optimized execution through JAX - Progress logging and convergence diagnostics - Scientifically validated (7/7 validation tests passed, T036-T041) - Serves as fallback when NLSQAdapter fails

Per-Angle Scaling Fix (v2.2): - Fixes silent optimization failures with per-angle parameters on large datasets - Applies angle-stratified chunking when: per_angle_scaling=True AND n_points>100k - Ensures every NLSQ chunk contains all phi angles → gradients always well-defined - <1% performance overhead (0.15s for 3M points) - Reference: ultra-think-20251106-012247

Production Status: - Production-ready with comprehensive error recovery - Scientifically validated (100% test pass rate) - Parameter recovery accuracy: 2-14% on core parameters - Sub-linear performance scaling with dataset size - Per-angle scaling compatible with large datasets (v2.2+)

References: - NLSQ Package: https://github.com/imewei/NLSQ - Validation: See tests/validation/test_scientific_validation.py (T036-T041) - Documentation: See CHANGELOG.md and CLAUDE.md for detailed status

homodyne.optimization.nlsq.wrapper.create_multistart_warmup_func(model_func, xdata, ydata, bounds=None, warmup_learning_rate=0.001, normalize=True, chunk_size=50_000)

Create a warmup-only fit function for multi-start Phase 1 strategy.

This function creates a warmup_fit_func compatible with the multi-start optimization module Phase 1 strategy. It uses the L-BFGS warmup phase from the NLSQ AdaptiveHybridStreamingOptimizer to quickly explore the parameter space without full Gauss-Newton refinement.

Parameters:

• model_func (Callable[..., ndarray]) – Model function with signature: func(x, *params) -> predictions

• xdata (ndarray) – Independent variable data

• ydata (ndarray) – Dependent variable data (observations)

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds as (lower, upper)

• warmup_learning_rate (float) – L-BFGS line search scale for warmup phase

• normalize (bool) – Whether to use parameter normalization (recommended for scale imbalance)

• chunk_size (int) – Points per chunk for streaming computation

Returns:

warmup_fit_func – Function with signature: (data, initial_params, n_iterations) -> SingleStartResult Compatible with run_multistart_nlsq() warmup_fit_func parameter.

Return type:

Callable[[dict[str, Any], ndarray, int], Any]

Raises:

RuntimeError – If AdaptiveHybridStreamingOptimizer is not available (NLSQ < 0.3.2)

Examples

>>> from homodyne.optimization.nlsq.wrapper import create_multistart_warmup_func
>>> from homodyne.optimization.nlsq.multistart import run_multistart_nlsq
>>>
>>> # Create warmup function
>>> warmup_func = create_multistart_warmup_func(
...     model_func=my_model,
...     xdata=x_data,
...     ydata=y_data,
...     bounds=(lower, upper),
... )
>>>
>>> # Use with multi-start
>>> result = run_multistart_nlsq(
...     data=my_data,
...     bounds=bounds,
...     config=config,
...     single_fit_func=full_fit_func,
...     warmup_fit_func=warmup_func,  # For Phase 1 strategy
... )

Notes

This function integrates with the Phase 1 multi-start strategy which: 1. Runs parallel L-BFGS warmup from multiple starting points 2. Selects the best warmup result 3. Performs full Gauss-Newton refinement from the best starting point

This approach is memory-efficient for very large datasets (>100M points) and provides good exploration of the parameter space.

See also

homodyne.optimization.nlsq.multistart.run_multistart_nlsq

Main multi-start function

homodyne.optimization.nlsq.multistart._run_phase1_strategy

Phase 1 strategy implementation

class homodyne.optimization.nlsq.wrapper.NLSQWrapper

Bases: NLSQAdapterBase

Adapter class for NLSQ package integration with homodyne optimization.

This class translates between homodyne’s optimization API and the NLSQ package’s curve_fit interface, handling: - Data format transformations - Parameter validation and bounds checking - Automatic strategy selection for large datasets - Hybrid error handling and recovery

Usage:

wrapper = NLSQWrapper(enable_large_dataset=True) result = wrapper.fit(data, config, initial_params, bounds, analysis_mode)

__init__(enable_large_dataset=True, enable_recovery=True, enable_numerical_validation=True, max_retries=2, fast_mode=False)

Initialize NLSQWrapper.

Parameters:

• enable_large_dataset (bool) – Use curve_fit_large for datasets >1M points

• enable_recovery (bool) – Enable automatic error recovery strategies

• enable_numerical_validation (bool) – Enable NaN/Inf validation at 3 critical points

• max_retries (int) – Maximum retry attempts per batch (default: 2)

• fast_mode (bool) – Disable non-essential checks for < 1% overhead (Task 5.5)

fit(data, config, initial_params=None, bounds=None, analysis_mode='static_isotropic', per_angle_scaling=True, diagnostics_enabled=False, shear_transforms=None, per_angle_scaling_initial=None)

Execute NLSQ optimization with automatic strategy selection and per-angle scaling.

Parameters:

• data (Any) – XPCS experimental data

• config (Any) – Configuration manager with optimization settings

• initial_params (ndarray | None) – Initial parameter guess (auto-loaded if None)

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds as (lower, upper) tuple

• analysis_mode (str) – ‘static_isotropic’ or ‘laminar_flow’

• per_angle_scaling (bool) – MUST be True. Per-angle contrast/offset parameters are physically correct as each scattering angle has different optical properties and detector responses. Legacy scalar mode (False) is no longer supported.

Return type:

OptimizationResult

Returns:

OptimizationResult with converged parameters and diagnostics

Raises:

ValueError – If bounds are invalid (lower > upper) or if per_angle_scaling=False

Key Features

• Automatic strategy selection based on dataset size

• Memory-aware chunking for large datasets

• JIT-compiled residual functions

• Stratified sampling for per-angle scaling

NLSQAdapter

Modern adapter for NLSQ v0.4+ CurveFit class with model caching and JIT support. This is the recommended path for new optimizations.

NLSQ Adapter using CurveFit class for homodyne optimization.

Role and When to Use (v2.11.0+)

NLSQAdapter (this module) is the recommended adapter for: - Standard optimizations (static_isotropic mode) - Small to medium datasets (< 10M points) - Multi-start optimization (model caching provides 3-5× speedup) - Performance-critical workflows requiring JIT compilation

Use NLSQWrapper instead for: - Complex optimizations requiring full anti-degeneracy integration - laminar_flow mode with many phi angles (> 6) - Large datasets (> 100M points) requiring streaming/chunking strategies - Custom transforms or advanced recovery mechanisms

Key Differences:

• Model caching: NLSQAdapter=Built-in, NLSQWrapper=None

• JIT compilation: NLSQAdapter=Auto, NLSQWrapper=Manual

• Workflow auto-select: NLSQAdapter=Via NLSQ, NLSQWrapper=Custom

• Anti-degeneracy layers: NLSQAdapter=Via fit(), NLSQWrapper=Full

• Recovery system: NLSQAdapter=NLSQ native, NLSQWrapper=3-attempt

• Streaming support: NLSQAdapter=Via NLSQ, NLSQWrapper=Full custom

Decision Guide:

1. If you need maximum speed for multi-start optimization: Use NLSQAdapter

2. If you need robust streaming for 100M+ points: Use NLSQWrapper

3. If you need full anti-degeneracy control: Use NLSQWrapper

4. Default recommendation for new code: Use NLSQAdapter (via use_adapter=True)

This module provides a modern adapter layer between homodyne’s optimization API and the NLSQ package’s CurveFit class, leveraging: - CurveFit class for JIT compilation caching - Model instance caching (WeakValueDictionary) for multi-start speedup - WorkflowSelector for automatic strategy selection - Built-in stability and recovery systems - Runtime fallback to NLSQWrapper on failure

This is the recommended integration path for NLSQ v0.4+ (homodyne v2.11.0+).

Key Features: - Model caching: 3-5× speedup for multi-start optimization - JIT compilation: 2-3× speedup for single fits - Automatic workflow selection based on dataset size and memory - Native NLSQ stability and recovery systems - Integration with homodyne’s anti-degeneracy defense system - Backward-compatible interface with NLSQWrapper.fit() - Automatic fallback to NLSQWrapper when adapter fails

Migration Guide: - Replace NLSQWrapper with NLSQAdapter - Set use_adapter=True in fit_nlsq_jax() (default in v2.11.0+) - Anti-degeneracy layers work unchanged

References: - NLSQ Package: https://github.com/imewei/NLSQ - Architecture: See CLAUDE.md for NLSQ integration details

class homodyne.optimization.nlsq.adapter.ModelCacheKey

Bases: object

Immutable key for model cache lookup.

Hashable tuple of (analysis_mode, phi_angles_tuple, q, per_angle_scaling). NumPy arrays converted to tuples for hashability.

analysis_mode

“static_isotropic” or “laminar_flow”

phi_angles

Unique phi angles (sorted) as tuple

q

Scattering wavevector magnitude

per_angle_scaling

Whether per-angle contrast/offset is used

analysis_mode: str

phi_angles: tuple[float, ...]

q: float

per_angle_scaling: bool

__init__(analysis_mode, phi_angles, q, per_angle_scaling)

class homodyne.optimization.nlsq.adapter.CachedModel

Bases: object

Cached model instance with JIT-compiled prediction function.

Stored in dict with LRU eviction - oldest entries removed when cache is full.

model

CombinedModel instance for computing g1/g2 values

model_func

Model prediction function (NumPy-compatible wrapper)

created_at

time.time() for diagnostics

n_hits

Cache hit counter for monitoring

model: Any

model_func: Callable[[ndarray, Any], ndarray]

created_at: float

n_hits: int = 0

__init__(model, model_func, created_at=<factory>, n_hits=0)

homodyne.optimization.nlsq.adapter.get_or_create_model(analysis_mode, phi_angles, q, per_angle_scaling=True, config=None, enable_jit=True)

Get cached model or create new one.

This function provides model instance caching to avoid redundant model creation during multi-start optimization. Expected 3-5× speedup.

Uses CombinedModel (not HomodyneModel) for simpler initialization. The model function closure captures the model and experimental setup.

Parameters:

• analysis_mode (str) – ‘static_isotropic’ or ‘laminar_flow’

• phi_angles (ndarray) – Unique phi angles in radians

• q (float) – Scattering wavevector magnitude

• per_angle_scaling (bool) – Whether per-angle contrast/offset is used

• config (dict[str, Any] | None) – Optional config dict for model initialization

• enable_jit (bool) – Whether to JIT-compile the model function

Returns:

• model: CombinedModel instance (cached or newly created)

• model_func: Prediction function (JIT-compiled if enable_jit=True)

• cache_hit: True if model was retrieved from cache

Return type:

tuple[Any, Callable[[ndarray, Any], ndarray], bool]

Raises:

ValueError – If analysis_mode is invalid, phi_angles is empty, or q <= 0

Example

>>> model, model_func, hit = get_or_create_model(
...     "laminar_flow",
...     np.array([0.0, 0.5, 1.0]),
...     0.001,
... )
>>> if hit:
...     logger.debug("Model cache hit")

homodyne.optimization.nlsq.adapter.clear_model_cache()

Clear all cached models.

Return type:

int

Returns:

Number of models removed from cache

Notes

Useful for testing or when configuration changes require fresh models.

homodyne.optimization.nlsq.adapter.get_cache_stats()

Get cache statistics.

Returns:

• “hits”: Cache hit count

• ”misses”: Cache miss count

• ”size”: Current cache size

Return type:

dict[str, int]

class homodyne.optimization.nlsq.adapter.AdapterConfig

Bases: object

Configuration for NLSQAdapter.

enable_cache

Enable model instance caching (new in v2.11.0)

enable_jit

Enable JIT compilation of model functions (new in v2.11.0)

enable_recovery

Enable NLSQ’s built-in recovery system

enable_stability

Enable NLSQ’s numerical stability guard

goal

Optimization goal (fast, robust, quality, memory_efficient)

workflow

Workflow tier override (auto, standard, streaming)

enable_cache: bool = True

enable_jit: bool = True

enable_recovery: bool = True

enable_stability: bool = True

goal: str = 'quality'

workflow: str = 'auto'

__init__(enable_cache=True, enable_jit=True, enable_recovery=True, enable_stability=True, goal='quality', workflow='auto')

class homodyne.optimization.nlsq.adapter.NLSQAdapter

Bases: NLSQAdapterBase

Adapter for NLSQ package using CurveFit class.

Uses NLSQ’s CurveFit for JIT caching and WorkflowSelector for automatic strategy selection. This is the modern integration path for NLSQ v0.4+ with improved performance and reliability.

Usage:

adapter = NLSQAdapter() result = adapter.fit(data, config, initial_params, bounds, analysis_mode)

Compared to NLSQWrapper:

• Uses CurveFit class for JIT compilation caching

• Leverages WorkflowSelector for auto strategy selection

• Delegates recovery to NLSQ’s built-in systems

• Simpler codebase with less custom logic

Note

Anti-degeneracy layers (hierarchical, shear_weighting, etc.) remain in homodyne as they are physics-specific to XPCS analysis.

__init__(config=None)

Initialize NLSQAdapter.

Parameters:

config (AdapterConfig | None) – Adapter configuration. If None, uses defaults.

Raises:

ImportError – If NLSQ CurveFit class is not available.

fit(data, config, initial_params=None, bounds=None, analysis_mode='static_isotropic', per_angle_scaling=True, diagnostics_enabled=False, shear_transforms=None, per_angle_scaling_initial=None, anti_degeneracy_controller=None)

Execute NLSQ optimization using CurveFit class.

This method provides the same interface as NLSQWrapper.fit() for backward compatibility while using NLSQ’s modern CurveFit class.

Parameters:

• data (Any) – XPCS experimental data

• config (Any) – Configuration manager with optimization settings

• initial_params (ndarray | None) – Initial parameter guess (required)

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds as (lower, upper) tuple

• analysis_mode (str) – ‘static_isotropic’ or ‘laminar_flow’

• per_angle_scaling (bool) – Must be True (per-angle is physically correct)

• diagnostics_enabled (bool) – Enable extended diagnostics

• shear_transforms (dict[str, Any] | None) – Shear parameter transformations

• per_angle_scaling_initial (dict[str, list[float]] | None) – Initial per-angle contrast/offset

• anti_degeneracy_controller (Any | None) – Anti-degeneracy controller (physics-specific)

Return type:

OptimizationResult

Returns:

OptimizationResult with converged parameters and diagnostics

Raises:

• ValueError – If bounds are invalid or per_angle_scaling=False

• ImportError – If NLSQ CurveFit is not available

is_available()

Check if NLSQ CurveFit is available.

Return type:

bool

property workflow_available: bool

Check if NLSQ WorkflowSelector is available.

homodyne.optimization.nlsq.adapter.get_adapter(config=None)

Factory function to get NLSQAdapter instance.

Parameters:

config (AdapterConfig | None) – Adapter configuration

Return type:

NLSQAdapter

Returns:

NLSQAdapter instance

Raises:

ImportError – If NLSQ CurveFit is not available

homodyne.optimization.nlsq.adapter.is_adapter_available()

Check if NLSQAdapter can be used.

Return type:

bool

Returns:

True if NLSQ CurveFit class is available

Key Classes

homodyne.optimization.nlsq.adapter.NLSQAdapter	Adapter for NLSQ package using CurveFit class.
homodyne.optimization.nlsq.adapter.AdapterConfig	Configuration for NLSQAdapter.
homodyne.optimization.nlsq.adapter.ModelCacheKey	Immutable key for model cache lookup.
homodyne.optimization.nlsq.adapter.CachedModel	Cached model instance with JIT-compiled prediction function.

Key Features

Model Caching (3-5× Multi-Start Speedup):

• Cached model instances avoid redundant model creation

• LRU eviction with 64-entry cache limit

• Cache hit/miss statistics for monitoring

JIT Compilation Flag:

• Signals intent for JIT optimization

• Underlying CombinedModel uses JAX internally

• Graceful fallback if JAX unavailable

Automatic Fallback:

• NLSQAdapter failures automatically retry with NLSQWrapper

• Logged warnings include original error for debugging

• Fallback metadata in device_info

Configuration

from homodyne.optimization.nlsq import AdapterConfig, NLSQAdapter

config = AdapterConfig(
    enable_cache=True,      # Model caching (default: True)
    enable_jit=True,        # JIT compilation (default: True)
    enable_recovery=True,   # NLSQ recovery system
    goal="quality",         # Optimization goal
)

adapter = NLSQAdapter(config)
result = adapter.fit(data, config, initial_params, bounds, analysis_mode)

Cache Management

from homodyne.optimization.nlsq import get_cache_stats, clear_model_cache

# View cache statistics
stats = get_cache_stats()
print(f"Hits: {stats['hits']}, Misses: {stats['misses']}")

# Clear cache (useful for testing)
n_cleared = clear_model_cache()

When to Use Which Adapter

Adapter	Use Case	Advantages
NLSQAdapter	Multi-start optimization, repeated fits	Model caching, modern API
NLSQWrapper	Complex workflows, anti-degeneracy	Full feature set, streaming support

Memory Management

Utilities for adaptive memory thresholds and streaming decisions.

Memory Management and Unified Strategy Selection for NLSQ Optimization.

Provides adaptive memory threshold detection and unified memory-based strategy selection for NLSQ optimization (v2.13.0+).

Key Features: - Cross-platform system memory detection (psutil + os.sysconf fallback) - Adaptive threshold calculation based on available memory - Unified memory-based strategy selection (NLSQStrategy, select_nlsq_strategy) - Environment variable override support (NLSQ_MEMORY_FRACTION) - Safe fraction clamping to prevent OOM or underutilization

Strategy Selection (v2.13.0+):

>>> from homodyne.optimization.nlsq.memory import select_nlsq_strategy
>>> decision = select_nlsq_strategy(n_points=100_000_000, n_params=53)
>>> print(decision.strategy.value)  # 'standard', 'out_of_core', or 'hybrid_streaming'

Memory Threshold:

>>> from homodyne.optimization.nlsq.memory import get_adaptive_memory_threshold
>>> threshold_gb, info = get_adaptive_memory_threshold()
>>> print(f"Threshold: {threshold_gb:.1f} GB")

homodyne.optimization.nlsq.memory.detect_total_system_memory()

Detect total system memory in bytes using multiple methods.

Returns:

Total system memory in bytes, or None if detection fails.

Return type:

float | None

Notes

Detection priority: 1. psutil.virtual_memory().total (preferred, cross-platform) 2. os.sysconf(‘SC_PAGE_SIZE’) * os.sysconf(‘SC_PHYS_PAGES’) (Linux fallback)

homodyne.optimization.nlsq.memory.get_adaptive_memory_threshold(memory_fraction=None)

Compute adaptive memory threshold based on system memory.

The memory threshold determines when NLSQ switches to streaming mode for memory-bounded optimization. Instead of a fixed 16 GB threshold, this function computes an adaptive threshold as a fraction of total system memory.

Parameters:

memory_fraction (float | None) – Fraction of total system memory to use as threshold (0.1 to 0.9). If None, uses: 1. Environment variable NLSQ_MEMORY_FRACTION (if set) 2. Default value of 0.75 (75% of total memory)

Return type:

tuple[float, dict[str, Any]]

Returns:

• threshold_gb (float) – Memory threshold in gigabytes.

• info (dict) – Diagnostic information with keys: - ‘total_memory_gb’: Detected total system memory (GB) - ‘memory_fraction’: Fraction used - ‘source’: How the fraction was determined (‘argument’, ‘env’, ‘default’) - ‘detection_method’: How memory was detected (‘psutil’, ‘sysconf’, ‘fallback’)

Notes

• If total memory cannot be detected, falls back to 16.0 GB with a warning.

• Memory fraction is clamped to [0.1, 0.9] for safety.

• Environment variable NLSQ_MEMORY_FRACTION can override the default.

Examples

>>> threshold_gb, info = get_adaptive_memory_threshold()
>>> pct = info['memory_fraction'] * 100
>>> tot = info['total_memory_gb']
>>> print(f"Threshold: {threshold_gb:.1f} GB ({pct:.0f}% of {tot:.1f} GB)")
Threshold: 24.0 GB (75% of 32.0 GB)

>>> # Override with specific fraction
>>> threshold_gb, _ = get_adaptive_memory_threshold(memory_fraction=0.5)

>>> # Override via environment variable
>>> import os
>>> os.environ["NLSQ_MEMORY_FRACTION"] = "0.6"
>>> threshold_gb, info = get_adaptive_memory_threshold()
>>> assert info['source'] == 'env'

homodyne.optimization.nlsq.memory.estimate_peak_memory_gb(n_points, n_params, bytes_per_element=8, jacobian_overhead=6.5)

Estimate peak memory usage for full Jacobian optimization.

Parameters:

• n_points (int) – Number of data points

• n_params (int) – Number of parameters

• bytes_per_element (int) – Bytes per float element (default: 8 for float64)

• jacobian_overhead (float) – Multiplicative factor accounting for: - Base Jacobian matrix (n_points × n_params) - Autodiff intermediate tensors (~2×) - Stratified array padding and copies (~1.5×) - JIT compilation intermediates (~1.5×) - Optimizer working memory (residuals, QR decomp, etc.) Default: 6.5 (empirically validated for 23M+ point datasets)

Returns:

Estimated peak memory in gigabytes

Return type:

float

class homodyne.optimization.nlsq.memory.NLSQStrategy

Bases: Enum

NLSQ optimization strategy based on memory constraints.

STANDARD = 'standard'

OUT_OF_CORE = 'out_of_core'

HYBRID_STREAMING = 'hybrid_streaming'

class homodyne.optimization.nlsq.memory.StrategyDecision

Bases: object

Result of unified memory-based strategy selection.

strategy

Selected optimization strategy

Type:

NLSQStrategy

threshold_gb

Memory threshold used for decision (GB)

Type:

float

index_memory_gb

Memory required for int64 index array (GB)

Type:

float

peak_memory_gb

Estimated peak memory for full Jacobian (GB)

Type:

float

reason

Human-readable explanation of decision

Type:

str

strategy: NLSQStrategy

threshold_gb: float

index_memory_gb: float

peak_memory_gb: float

reason: str

__init__(strategy, threshold_gb, index_memory_gb, peak_memory_gb, reason)

homodyne.optimization.nlsq.memory.select_nlsq_strategy(n_points, n_params, memory_fraction=DEFAULT_MEMORY_FRACTION)

Unified memory-based NLSQ strategy selection.

Implements a pure memory-based decision tree:

1. If index array > threshold → HYBRID_STREAMING (extreme scale)

2. Elif peak memory > threshold → OUT_OF_CORE (large scale)

3. Else → STANDARD (in-memory)

Parameters:

• n_points (int) – Number of data points

• n_params (int) – Number of optimization parameters

• memory_fraction (float) – Fraction of system RAM to use as threshold (default: 0.75)

Returns:

Decision with strategy, metrics, and rationale

Return type:

StrategyDecision

Examples

>>> decision = select_nlsq_strategy(100_000_000, 53)
>>> print(decision.strategy.value)
'out_of_core'
>>> print(decision.reason)
'Peak memory (12.8 GB) exceeds threshold (24.0 GB)'

Parameter Utilities

Helper functions for parameter handling and per-angle initialization.

Parameter Utilities for NLSQ Optimization.

Provides utility functions for parameter handling, labeling, status classification, and per-angle initialization in NLSQ optimization.

Key Functions: - build_parameter_labels: Create parameter labels with per-angle support - classify_parameter_status: Identify parameters at bounds - sample_xdata: Subsample x-data for diagnostic computations - compute_consistent_per_angle_init: Initialize per-angle params consistently - compute_jacobian_stats: Compute Jacobian-based statistics

homodyne.optimization.nlsq.parameter_utils.build_parameter_labels(per_angle_scaling, n_phi, physical_param_names)

Build parameter labels including per-angle scaling parameters.

Parameters:

• per_angle_scaling (bool) – Whether per-angle contrast/offset are used

• n_phi (int) – Number of phi angles

• physical_param_names (list[str]) – Names of physical parameters

Returns:

Full list of parameter labels

Return type:

list[str]

homodyne.optimization.nlsq.parameter_utils.classify_parameter_status(values, lower, upper, atol=1e-9)

Classify parameters as active or at bounds.

Parameters:

• values (ndarray) – Current parameter values

• lower (ndarray | None) – Lower bounds

• upper (ndarray | None) – Upper bounds

• atol (float) – Absolute tolerance for bound comparison

Returns:

Status for each parameter: ‘active’, ‘at_lower_bound’, or ‘at_upper_bound’

Return type:

list[str]

homodyne.optimization.nlsq.parameter_utils.sample_xdata(xdata, max_points)

Subsample x-data for diagnostic computations.

Parameters:

• xdata (ndarray) – Input data

• max_points (int) – Maximum number of points to return

Returns:

Subsampled data

Return type:

ndarray

homodyne.optimization.nlsq.parameter_utils.compute_jacobian_stats(residual_fn, x_subset, params, scaling_factor)

Compute Jacobian statistics for diagnostics.

Parameters:

• residual_fn (Callable[..., Any]) – Residual function

• x_subset (ndarray) – Subset of x data for computation

• params (ndarray) – Current parameters

• scaling_factor (float) – Scaling factor for statistics

Returns:

(J^T J matrix, column norms) or (None, None) on failure

Return type:

tuple[ndarray | None, ndarray | None]

homodyne.optimization.nlsq.parameter_utils.compute_consistent_per_angle_init(stratified_data, physical_params, physical_param_names, default_contrast=0.5, default_offset=1.0, logger=None)

Compute per-angle contrast/offset consistent with initial physical parameters.

This function solves a critical initialization problem in laminar_flow mode: when physical shear parameters (gamma_dot_t0) are nonzero, the model predicts DIFFERENT g2 values at different angles. If per-angle contrast/offset are initialized uniformly, large initial residuals can cause the optimizer to incorrectly reduce gamma_dot_t0 to zero.

Instead, we compute per-angle contrast/offset by fitting:

g2_data[angle] ≈ offset[angle] + contrast[angle] × g1_model²[angle]

where g1_model is computed using the initial physical parameters.

Parameters:

• stratified_data (Any) – Data containing per-angle g2, phi, t1, t2 arrays

• physical_params (ndarray) – Initial physical parameters [D0, alpha, D_offset, (gamma_dot_t0, beta, gamma_dot_t_offset, phi0)]

• physical_param_names (list[str]) – Names of physical parameters to determine analysis mode

• default_contrast (float) – Default contrast value if fitting fails

• default_offset (float) – Default offset value if fitting fails

• logger (Any) – Logger for diagnostic messages

Return type:

tuple[ndarray, ndarray]

Returns:

• contrast_per_angle (np.ndarray) – Per-angle contrast values consistent with physical params

• offset_per_angle (np.ndarray) – Per-angle offset values consistent with physical params

homodyne.optimization.nlsq.parameter_utils.compute_quantile_per_angle_scaling(stratified_data, contrast_bounds=(0.0, 1.0), offset_bounds=(0.5, 1.5), lag_floor_quantile=0.80, lag_ceiling_quantile=0.20, value_quantile_low=0.10, value_quantile_high=0.90, logger=None)

Estimate per-angle contrast/offset from quantiles of c2_experimental values.

This function uses physics-informed quantile analysis to estimate contrast and offset for each phi angle independently. Unlike least-squares fitting, this approach does not require a model and directly extracts scaling from the data.

Physics basis:

C2 = contrast × g1² + offset

• At large time lags, g1² → 0, so C2 → offset (the “floor”)

• At small time lags, g1² ≈ 1, so C2 ≈ contrast + offset (the “ceiling”)

Parameters:

• stratified_data (Any) – Data containing per-angle g2_flat, phi_flat, t1_flat, t2_flat arrays.

• contrast_bounds (tuple[float, float]) – Valid bounds for contrast parameter.

• offset_bounds (tuple[float, float]) – Valid bounds for offset parameter.

• lag_floor_quantile (float) – Quantile threshold for “large lag” region (default: 0.80 = top 20% of lags).

• lag_ceiling_quantile (float) – Quantile threshold for “small lag” region (default: 0.20 = bottom 20% of lags).

• value_quantile_low (float) – Quantile for robust floor estimation (default: 0.10).

• value_quantile_high (float) – Quantile for robust ceiling estimation (default: 0.90).

• logger (Any) – Logger for diagnostic messages.

Return type:

tuple[ndarray, ndarray]

Returns:

• contrast_per_angle (np.ndarray) – Per-angle contrast values from quantile estimation.

• offset_per_angle (np.ndarray) – Per-angle offset values from quantile estimation.

Notes

The estimation is robust to outliers by using quantiles instead of min/max. The lag-based segmentation ensures we sample from appropriate regions of the correlation decay curve.

This function is designed for the “constant” mode in anti-degeneracy defense, where per-angle contrast/offset are estimated once and treated as fixed parameters during optimization.

Parameter Index Mapper

Single source of truth for parameter indices across all anti-degeneracy modes.

Centralized index mapping for anti-degeneracy layers.

This module provides the ParameterIndexMapper class which ensures consistent index ranges regardless of whether Fourier reparameterization is active. This is the single source of truth for parameter group boundaries.

Created: 2025-12-31 Feature: 001-fix-nlsq-anti-degeneracy

class homodyne.optimization.nlsq.parameter_index_mapper.ParameterIndexMapper

Bases: object

Centralized index mapping for anti-degeneracy layers.

Provides consistent index ranges regardless of whether Fourier reparameterization or constant scaling is active. This class is the single source of truth for parameter group boundaries.

Parameters:

• n_phi (int) – Number of unique phi angles.

• n_physical (int) – Number of physical parameters (typically 7 for laminar_flow mode).

• fourier (FourierReparameterizer | None) – Reference to Fourier reparameterizer if Layer 1 is active.

• use_constant (bool) – Whether constant scaling mode is active (single contrast/offset shared across all angles).

n_per_angle_total

Total number of per-angle parameters (Fourier coefficients, raw, or 2).

Type:

int

n_per_group

Number of parameters per group (contrast or offset).

Type:

int

use_fourier

Whether Fourier reparameterization is active.

Type:

bool

use_constant

Whether constant scaling mode is active.

Type:

bool

total_params

Total number of parameters.

Type:

int

mode_name

Human-readable name of current mode (“constant”, “fourier”, or “individual”).

Type:

str

Examples

>>> # Constant mode (23 phi angles)
>>> mapper = ParameterIndexMapper(n_phi=23, n_physical=7, use_constant=True)
>>> mapper.get_group_indices()
[(0, 1), (1, 2)]
>>> mapper.n_per_angle_total
2
>>> mapper.mode_name
'constant'

>>> # Non-Fourier mode (23 phi angles)
>>> mapper = ParameterIndexMapper(n_phi=23, n_physical=7, fourier=None)
>>> mapper.get_group_indices()
[(0, 23), (23, 46)]
>>> mapper.n_per_angle_total
46

>>> # Fourier mode (23 phi angles, order=2)
>>> mapper = ParameterIndexMapper(n_phi=23, n_physical=7, fourier=fourier_obj)
>>> mapper.get_group_indices()
[(0, 5), (5, 10)]
>>> mapper.n_per_angle_total
10

n_phi: int

n_physical: int

fourier: FourierReparameterizer | None = None

use_constant: bool = False

__post_init__()

Validate inputs and cache computed values.

Return type:

None

property use_fourier: bool

Check if Fourier reparameterization is active.

property n_per_group: int

Get number of parameters per group (contrast or offset).

Returns:

1 for constant mode, n_coeffs for Fourier, n_phi for individual.

Return type:

int

property mode_name: str

Get human-readable name of current mode.

Returns:

“constant”, “fourier”, or “individual”

Return type:

str

property n_per_angle_total: int

Get total number of per-angle parameters (scaling params).

property total_params: int

Get total number of parameters.

get_group_indices()

Get (start, end) tuples for contrast and offset parameter groups.

Returns:

Two tuples: [(contrast_start, contrast_end), (offset_start, offset_end)]

Return type:

list[tuple[int, int]]

Notes

• Contrast group: indices [0, n_per_group)

• Offset group: indices [n_per_group, 2*n_per_group)

get_physical_indices()

Get indices of physical parameters.

Returns:

Indices of physical parameters in the full parameter vector.

Return type:

list[int]

get_per_angle_indices()

Get indices of all per-angle parameters.

Returns:

Indices of per-angle parameters (contrast + offset).

Return type:

list[int]

validate_indices(params)

Validate that group indices are within parameter vector bounds.

Parameters:

params (ndarray) – Full parameter vector.

Returns:

True if all indices are valid, False otherwise.

Return type:

bool

Raises:

ValueError – If indices are out of bounds (with descriptive message).

get_diagnostics()

Get diagnostic information for logging.

Returns:

Diagnostic information including mode, counts, and indices.

Return type:

dict

get_covariance_slice_indices()

Get slice indices for covariance matrix transformation.

Returns slices for extracting per-angle and physical parameter blocks from a covariance matrix.

Returns:

(per_angle_slice, physical_slice) for indexing covariance matrices.

Return type:

tuple[slice, slice]

__init__(n_phi, n_physical, fourier=None, use_constant=False)

Key Classes

homodyne.optimization.nlsq.parameter_index_mapper.ParameterIndexMapper

Centralized index mapping for anti-degeneracy layers.

Usage Example

from homodyne.optimization.nlsq.parameter_index_mapper import ParameterIndexMapper

# Constant mode (23 phi angles, 7 physical params)
mapper = ParameterIndexMapper(n_phi=23, n_physical=7, use_constant=True)
print(mapper.mode_name)           # "constant"
print(mapper.n_per_angle_total)   # 2 (single contrast + offset, shared)
print(mapper.total_params)        # 9 (2 + 7)

# Fourier mode (order=2)
mapper = ParameterIndexMapper(n_phi=23, n_physical=7, use_fourier=True, fourier_order=2)
print(mapper.mode_name)           # "fourier"
print(mapper.n_per_angle_total)   # 10 (5 contrast + 5 offset coefficients)
print(mapper.total_params)        # 17 (10 + 7)

Jacobian Utilities

Jacobian computation utilities for convergence diagnostics.

Jacobian computation utilities for NLSQ optimization.

This module extracts Jacobian-related functions from nlsq_wrapper.py to reduce file size and improve maintainability.

Extracted from nlsq_wrapper.py as part of technical debt remediation (Dec 2025).

homodyne.optimization.nlsq.jacobian.compute_jacobian_stats(residual_fn, x_subset, params, scaling_factor)

Compute Jacobian statistics for convergence diagnostics.

Computes the Jacobian matrix and derives: - JTJ (Jacobian transpose times Jacobian) for Hessian approximation - Column norms for parameter sensitivity analysis

Parameters:

• residual_fn (Callable[..., Any]) – Residual function to differentiate.

• x_subset (ndarray) – Subset of x data for Jacobian computation.

• params (ndarray) – Current parameter values.

• scaling_factor (float) – Scaling factor for JTJ computation.

Returns:

(JTJ matrix, column norms) or (None, None) on failure.

Return type:

tuple[ndarray | None, ndarray | None]

homodyne.optimization.nlsq.jacobian.compute_jacobian_condition_number(residual_fn, x_subset, params)

Compute condition number of Jacobian matrix.

The condition number indicates how sensitive the optimization is to parameter perturbations. High values (>1e6) suggest ill-conditioning.

Parameters:

• residual_fn (Callable[..., Any]) – Residual function to differentiate.

• x_subset (ndarray) – Subset of x data for Jacobian computation.

• params (ndarray) – Current parameter values.

Returns:

Condition number or None on failure.

Return type:

float | None

homodyne.optimization.nlsq.jacobian.analyze_parameter_sensitivity(residual_fn, x_subset, params, param_names)

Analyze parameter sensitivity from Jacobian column norms.

Higher column norms indicate parameters that have more influence on the residuals.

Parameters:

• residual_fn (Callable[..., Any]) – Residual function to differentiate.

• x_subset (ndarray) – Subset of x data for Jacobian computation.

• params (ndarray) – Current parameter values.

• param_names (list[str]) – Parameter names for labeling.

Returns:

Mapping from parameter name to sensitivity (normalized 0-1).

Return type:

dict[str, float]

homodyne.optimization.nlsq.jacobian.estimate_gradient_noise(residual_fn, x_subset, params, n_samples=5, perturbation=1e-6, seed=42)

Estimate gradient noise from multiple Jacobian computations.

Computes Jacobian multiple times with small perturbations to estimate numerical noise in gradient computation.

Parameters:

• residual_fn (Callable[..., Any]) – Residual function to differentiate.

• x_subset (ndarray) – Subset of x data for Jacobian computation.

• params (ndarray) – Current parameter values.

• n_samples (int) – Number of perturbed samples.

• perturbation (float) – Relative perturbation size.

Returns:

Estimated gradient noise (coefficient of variation) or None on failure.

Return type:

float | None

Progress Tracking

Progress bar and logging callbacks for NLSQ optimization.

Progress bar and logging callbacks for NLSQ optimization.

This module provides progress tracking for NLSQ fitting operations, integrating with the NLSQ package’s callback system.

Features: - tqdm progress bar for fitting operations - Iteration logging with configurable interval - Multi-start progress tracking - Streaming optimization progress

Part of homodyne v2.7.0 architecture.

class homodyne.optimization.nlsq.progress.ProgressConfig

Bases: object

Configuration for progress tracking.

enable_progress_bar

Whether to show tqdm progress bar.

Type:

bool

verbose

Verbosity level: 0=quiet, 1=normal, 2=detailed.

Type:

int

log_interval

Log every N iterations when verbose >= 2.

Type:

int

max_nfev

Maximum function evaluations (for progress bar total).

Type:

int

description

Description for progress bar.

Type:

str

enable_progress_bar: bool = True

verbose: int = 1

log_interval: int = 10

max_nfev: int = 1000

description: str = 'NLSQ Fitting'

classmethod from_nlsq_config(nlsq_config, max_nfev=None, description='NLSQ Fitting')

Create ProgressConfig from NLSQConfig.

Parameters:

• nlsq_config (NLSQConfig) – NLSQ configuration object.

• max_nfev (int | None) – Max function evaluations. Uses nlsq_config.max_iterations if None.

• description (str) – Description for progress bar.

Returns:

Progress configuration.

Return type:

ProgressConfig

__init__(enable_progress_bar=True, verbose=1, log_interval=10, max_nfev=1000, description='NLSQ Fitting')

class homodyne.optimization.nlsq.progress.HomodyneIterationLogger

Bases: object

Iteration logger that integrates with homodyne’s logging system.

Logs optimization progress at configurable intervals using the homodyne logging infrastructure.

Parameters:

• verbose (int) – Verbosity level: 0=quiet, 1=normal (milestones), 2=detailed.

• log_interval (int) – Log every N iterations when verbose >= 2.

• logger_instance (Logger | None) – Logger to use. Defaults to module logger.

__init__(verbose=1, log_interval=10, logger_instance=None)

__call__(iteration, cost, params, info)

Log iteration information based on verbosity settings.

Return type:

None

close()

Log final summary.

Return type:

None

homodyne.optimization.nlsq.progress.create_progress_callback(config=None, enable_progress_bar=True, verbose=1, log_interval=10, max_nfev=1000, description='NLSQ Fitting')

Create progress callback chain for NLSQ optimization.

Creates a callback chain with optional progress bar and iteration logger.

Parameters:

• config (ProgressConfig | None) – Progress configuration. If provided, overrides other parameters.

• enable_progress_bar (bool) – Whether to show tqdm progress bar.

• verbose (int) – Verbosity level: 0=quiet, 1=normal, 2=detailed.

• log_interval (int) – Log every N iterations when verbose >= 2.

• max_nfev (int) – Maximum function evaluations for progress bar.

• description (str) – Description for progress bar.

Returns:

(callback, iteration_logger) - callback for NLSQ, logger for manual close. Returns (None, None) if no callbacks are needed.

Return type:

tuple[CallbackBase | None, HomodyneIterationLogger | None]

class homodyne.optimization.nlsq.progress.MultiStartProgressTracker

Bases: object

Progress tracker for multi-start optimization.

Provides a progress bar and logging for multi-start optimization, tracking the progress of multiple starting points.

Parameters:

• n_starts (int) – Total number of starting points.

• enable_progress_bar (bool) – Whether to show tqdm progress bar.

• verbose (int) – Verbosity level.

• description (str) – Description for progress bar.

__init__(n_starts, enable_progress_bar=True, verbose=1, description='Multi-start NLSQ')

update(start_idx, success, chi_squared, message='', wall_time=None)

Update progress after a single start completes.

Parameters:

• start_idx (int) – Index of the completed starting point.

• success (bool) – Whether optimization was successful.

• chi_squared (float) – Final chi-squared value.

• message (str) – Status message.

• wall_time (float | None) – Time taken for this optimization in seconds.

Return type:

None

close()

Close progress bar and log summary.

Return type:

None

__enter__()

Context manager entry.

Return type:

MultiStartProgressTracker

__exit__(exc_type, exc_val, exc_tb)

Context manager exit.

Return type:

bool

homodyne.optimization.nlsq.progress.create_streaming_progress_callback(n_total_points, batch_size, max_epochs, enable_progress_bar=True, verbose=1)

Create a progress callback for streaming optimization.

Parameters:

• n_total_points (int) – Total number of data points.

• batch_size (int) – Batch size for streaming.

• max_epochs (int) – Maximum number of epochs.

• enable_progress_bar (bool) – Whether to show progress bar.

• verbose (int) – Verbosity level.

Returns:

Callback function for streaming optimizer, or None if not needed.

Return type:

Callable[[int, ndarray, float], bool] | None

Key Classes

homodyne.optimization.nlsq.progress.ProgressConfig

Configuration for progress tracking.

Parameter Transforms

Parameter transformation utilities and name normalization.

Parameter transformation utilities for NLSQ optimization.

This module extracts shear transform logic from nlsq_wrapper.py to reduce file size and improve maintainability.

Extracted from nlsq_wrapper.py as part of technical debt remediation (Dec 2025).

homodyne.optimization.nlsq.transforms.normalize_param_key(name)

Normalize parameter name using canonical aliases.

Parameters:

name (str | None) – Parameter name to normalize.

Returns:

Canonical parameter name.

Return type:

str

homodyne.optimization.nlsq.transforms.normalize_x_scale_map(raw_map)

Normalize parameter scaling map.

Parameters:

raw_map (Any) – Raw scaling map (dict or other).

Returns:

Normalized scaling map with canonical keys.

Return type:

dict[str, float]

homodyne.optimization.nlsq.transforms.build_per_parameter_x_scale(per_angle_scaling, n_angles, physical_param_names, analysis_mode, override_map)

Build per-parameter scale array for optimization.

Parameters:

• per_angle_scaling (bool) – Whether per-angle scaling is enabled.

• n_angles (int) – Number of phi angles.

• physical_param_names (list[str]) – List of physical parameter names.

• analysis_mode (str) – Analysis mode (“static” or “laminar_flow”).

• override_map (dict[str, float]) – User overrides for parameter scales.

Returns:

Scale array or None if all scales are 1.0.

Return type:

ndarray | None

homodyne.optimization.nlsq.transforms.format_x_scale_for_log(value)

Format x_scale value for logging.

Parameters:

value (Any) – Scale value to format.

Returns:

Formatted string.

Return type:

str

homodyne.optimization.nlsq.transforms.parse_shear_transform_config(config)

Parse shear transform configuration.

Parameters:

config (Any | None) – Configuration dict or None.

Returns:

Parsed configuration with defaults.

Return type:

dict[str, Any]

homodyne.optimization.nlsq.transforms.build_physical_index_map(per_angle_scaling, n_angles, physical_param_names)

Build mapping from parameter names to indices.

Parameters:

• per_angle_scaling (bool) – Whether per-angle scaling is enabled.

• n_angles (int) – Number of phi angles.

• physical_param_names (list[str]) – List of physical parameter names.

Returns:

Mapping from parameter name to index in parameter vector.

Return type:

dict[str, int]

homodyne.optimization.nlsq.transforms.apply_forward_shear_transforms_to_vector(params, index_map, transform_cfg)

Apply forward shear transforms to parameter vector.

Transforms parameters from physical space to solver space: - gamma_dot_t0 -> log(gamma_dot_t0) if enable_gamma_dot_log - beta -> beta - beta_reference if enable_beta_centering

Parameters:

• params (ndarray) – Parameter vector in physical space.

• index_map (dict[str, int]) – Mapping from parameter names to indices.

• transform_cfg (dict[str, Any]) – Transform configuration.

Returns:

(transformed_params, transform_state)

Return type:

tuple[ndarray, dict[str, Any]]

homodyne.optimization.nlsq.transforms.apply_forward_shear_transforms_to_bounds(bounds, state)

Apply forward shear transforms to parameter bounds.

Parameters:

• bounds (tuple[ndarray, ndarray] | None) – (lower, upper) bounds in physical space.

• state (dict[str, Any]) – Transform state from apply_forward_shear_transforms_to_vector.

Returns:

Transformed bounds or None.

Return type:

tuple[ndarray, ndarray] | None

homodyne.optimization.nlsq.transforms.apply_inverse_shear_transforms_to_vector(params, state)

Apply inverse shear transforms to parameter vector.

Transforms parameters from solver space back to physical space.

Parameters:

• params (ndarray) – Parameter vector in solver space.

• state (dict[str, Any] | None) – Transform state from apply_forward_shear_transforms_to_vector.

Returns:

Parameter vector in physical space.

Return type:

ndarray

homodyne.optimization.nlsq.transforms.adjust_covariance_for_transforms(covariance, transformed_params, physical_params, state)

Adjust covariance matrix for parameter transforms.

Parameters:

• covariance (ndarray) – Covariance matrix in solver space.

• transformed_params (ndarray) – Parameters in solver space.

• physical_params (ndarray) – Parameters in physical space.

• state (dict[str, Any] | None) – Transform state.

Returns:

Covariance matrix in physical space.

Return type:

ndarray

homodyne.optimization.nlsq.transforms.wrap_model_function_with_transforms(model_fn, state)

Wrap model function to apply inverse transforms to parameters.

Parameters:

• model_fn (Any) – Original model function.

• state (dict[str, Any] | None) – Transform state.

Returns:

Wrapped model function (or original if no transforms).

Return type:

Any

homodyne.optimization.nlsq.transforms.wrap_stratified_function_with_transforms(residual_fn, state)

Wrap stratified residual function with transforms.

Parameters:

• residual_fn (Any) – Original stratified residual function.

• state (dict[str, Any] | None) – Transform state.

Returns:

Wrapped function (or original if no transforms).

Return type:

Any

Results

NLSQ optimization result classes.

This module extracts result dataclasses from nlsq_wrapper.py to reduce file size and improve maintainability.

Extracted from nlsq_wrapper.py as part of technical debt remediation (Dec 2025).

class homodyne.optimization.nlsq.results.FunctionEvaluationCounter

Bases: object

Wraps a callable and counts invocations.

Useful for tracking the number of function evaluations during optimization.

fn: Callable[[...], Any]

count: int = 0

__call__(*args, **kwargs)

Call the wrapped function and increment count.

__init__(fn, count=0)

class homodyne.optimization.nlsq.results.OptimizationResult

Bases: object

Complete optimization result with fit quality metrics and diagnostics.

parameters

Converged parameter values.

Type:

np.ndarray

uncertainties

Standard deviations from covariance matrix diagonal.

Type:

np.ndarray

covariance

Full parameter covariance matrix.

Type:

np.ndarray

chi_squared

Sum of squared residuals.

Type:

float

reduced_chi_squared

chi_squared / (n_data - n_params).

Type:

float

convergence_status

‘converged’, ‘max_iter’, or ‘failed’.

Type:

str

iterations

Number of optimization iterations.

Type:

int

execution_time

Wall-clock execution time in seconds.

Type:

float

device_info

Device used for computation (CPU details).

Type:

dict[str, Any]

recovery_actions

List of error recovery actions taken.

Type:

list[str]

quality_flag

‘good’, ‘marginal’, or ‘poor’.

Type:

str

streaming_diagnostics

Enhanced diagnostics for streaming optimization.

Type:

dict[str, Any] | None

stratification_diagnostics

Diagnostics for angle-stratified chunking.

Type:

StratificationDiagnostics | None

nlsq_diagnostics

Additional NLSQ-specific diagnostics.

Type:

dict[str, Any] | None

parameters: ndarray

uncertainties: ndarray

covariance: ndarray

chi_squared: float

reduced_chi_squared: float

convergence_status: str

iterations: int

execution_time: float

device_info: dict[str, Any]

recovery_actions: list[str]

quality_flag: str = 'good'

streaming_diagnostics: dict[str, Any] | None = None

stratification_diagnostics: StratificationDiagnostics | None = None

nlsq_diagnostics: dict[str, Any] | None = None

sigma_is_default: bool = False

property success: bool

Return True if optimization converged (backward compatibility).

property message: str

Return descriptive message about optimization outcome.

__init__(parameters, uncertainties, covariance, chi_squared, reduced_chi_squared, convergence_status, iterations, execution_time, device_info, recovery_actions=<factory>, quality_flag='good', streaming_diagnostics=None, stratification_diagnostics=None, nlsq_diagnostics=None, sigma_is_default=False)

class homodyne.optimization.nlsq.results.FallbackInfo

Bases: object

Tracks fallback from NLSQAdapter to NLSQWrapper.

Included in OptimizationResult.device_info when fallback occurs.

fallback_occurred

True if fallback was triggered

adapter_used

“NLSQAdapter” or “NLSQWrapper”

adapter_error

Error message if adapter failed (None if succeeded)

wrapper_error

Error message if wrapper also failed (None otherwise)

States:

• NLSQAdapter + fallback_occurred=False + adapter_error=None: Adapter succeeded

• NLSQWrapper + fallback_occurred=True + adapter_error=”…”: Fallback succeeded

• NLSQWrapper + fallback_occurred=True + adapter_error=”…” + wrapper_error=”…”: Both failed

fallback_occurred: bool

adapter_used: str

adapter_error: str | None = None

wrapper_error: str | None = None

to_dict()

Convert to dict for inclusion in device_info.

Return type:

dict[str, Any]

__init__(fallback_occurred, adapter_used, adapter_error=None, wrapper_error=None)

class homodyne.optimization.nlsq.results.UseSequentialOptimization

Bases: object

Marker indicating sequential per-angle optimization should be used.

This is returned by _apply_stratification_if_needed when conditions require sequential per-angle optimization as a fallback strategy.

data

Original XPCS data object.

Type:

Any

reason

Why sequential optimization is needed.

Type:

str

data: Any

reason: str

__init__(data, reason)

Data Preparation

Data preparation utilities for NLSQ optimization.

Data Preparation Utilities for NLSQ Optimization.

This module provides data preparation functions extracted from wrapper.py to improve code organization and reduce complexity.

Extracted from wrapper.py as part of refactoring (Dec 2025).

class homodyne.optimization.nlsq.data_prep.PreparedData

Bases: object

Container for prepared optimization data.

xdata

Flattened independent variable data

ydata

Flattened dependent variable data (observations)

n_data

Total number of data points

n_phi

Number of unique phi angles

phi_unique

Unique phi angle values

xdata: ndarray

ydata: ndarray

n_data: int

n_phi: int

phi_unique: ndarray

__init__(xdata, ydata, n_data, n_phi, phi_unique)

class homodyne.optimization.nlsq.data_prep.ExpandedParameters

Bases: object

Container for expanded per-angle parameters.

params

Expanded parameter array

bounds

Expanded bounds tuple (lower, upper)

n_params

Total number of parameters

n_physical

Number of physical parameters

n_angles

Number of angles

params: ndarray

bounds: tuple[ndarray, ndarray] | None

n_params: int

n_physical: int

n_angles: int

__init__(params, bounds, n_params, n_physical, n_angles)

homodyne.optimization.nlsq.data_prep.expand_per_angle_parameters(compact_params, compact_bounds, n_angles, n_physical, logger=None)

Expand compact parameters to per-angle format.

When per_angle_scaling=True with N angles, parameters are structured as: - N contrast parameters (one per angle) - N offset parameters (one per angle) - n_physical physical parameters

Input (compact): [contrast, offset, physical_params…] Output (expanded): [c0, c1, …, cN-1, o0, o1, …, oN-1, physical_params…]

Parameters:

• compact_params (ndarray) – Compact parameter array (n_physical + 2 elements)

• compact_bounds (tuple[ndarray, ndarray] | None) – Compact bounds tuple or None

• n_angles (int) – Number of phi angles

• n_physical (int) – Number of physical parameters

• logger (Any) – Optional logger for diagnostics

Return type:

ExpandedParameters

Returns:

ExpandedParameters with per-angle parameters and bounds

Raises:

ValueError – If parameter count doesn’t match expected

homodyne.optimization.nlsq.data_prep.validate_bounds(bounds, n_params, logger=None)

Validate parameter bounds.

Parameters:

• bounds (tuple[ndarray, ndarray] | None) – Bounds tuple (lower, upper) or None

• n_params (int) – Expected number of parameters

• logger (Any) – Optional logger for diagnostics

Return type:

tuple[ndarray, ndarray] | None

Returns:

Validated bounds or None

Raises:

ValueError – If bounds are invalid

homodyne.optimization.nlsq.data_prep.validate_initial_params(params, bounds, logger=None)

Validate and clip initial parameters to bounds.

Parameters:

• params (ndarray) – Initial parameter guess

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds or None

• logger (Any) – Optional logger for diagnostics

Return type:

ndarray

Returns:

Validated parameters (clipped to bounds if needed)

homodyne.optimization.nlsq.data_prep.convert_bounds_to_nlsq_format(bounds)

Convert bounds to NLSQ-compatible format.

NLSQ expects bounds as (lower_array, upper_array) with float64 dtype.

Parameters:

bounds (tuple[ndarray, ndarray] | tuple[list, list] | None) – Input bounds in various formats

Return type:

tuple[ndarray, ndarray] | None

Returns:

Bounds as (lower, upper) numpy arrays or None

homodyne.optimization.nlsq.data_prep.build_parameter_labels(per_angle_scaling, n_phi, physical_param_names)

Build human-readable parameter labels.

Parameters:

• per_angle_scaling (bool) – Whether per-angle scaling is enabled

• n_phi (int) – Number of phi angles

• physical_param_names (list[str]) – Names of physical parameters

Return type:

list[str]

Returns:

List of parameter labels

homodyne.optimization.nlsq.data_prep.classify_parameter_status(values, lower, upper, atol=1e-9)

Classify parameter status relative to bounds.

Parameters:

• values (ndarray) – Parameter values

• lower (ndarray | None) – Lower bounds or None

• upper (ndarray | None) – Upper bounds or None

• atol (float) – Absolute tolerance for bound comparison

Returns:

‘active’, ‘at_lower_bound’, ‘at_upper_bound’

Return type:

list[str]

Key Classes

homodyne.optimization.nlsq.data_prep.PreparedData	Container for prepared optimization data.
homodyne.optimization.nlsq.data_prep.ExpandedParameters	Container for expanded per-angle parameters.

Fit Computation

Utilities for computing theoretical fits from NLSQ results.

Fit Computation Utilities for NLSQ Results.

This module provides functions for computing theoretical fits from NLSQ optimization results. Extracted from cli/commands.py for better organization.

Extracted from cli/commands.py as part of refactoring (Dec 2025).

homodyne.optimization.nlsq.fit_computation.compute_g2_batch(physical_params, t1, t2, phi_angles, q, L, dt, contrast=1.0, offset=1.0)

Compute g2 for all phi angles in a single vectorized operation.

Performance Optimization (Spec 006 - FR-007): Uses jax.vmap to compute g2 for all angles in parallel instead of sequential Python loop. Expected speedup: 10-20x for post-fitting.

Parameters:

• physical_params (Array) – Physical parameters array

• t1 (Array) – t1 time values, shape (n_t1,)

• t2 (Array) – t2 time values, shape (n_t2,)

• phi_angles (Array) – Phi angles in radians, shape (n_phi,)

• q (float) – Wave vector magnitude

• L (float) – Sample-to-detector distance

• dt (float) – Time step

• contrast (float) – Contrast parameter (default 1.0 for raw computation)

• offset (float) – Offset parameter (default 1.0 for raw computation)

Returns:

g2 values, shape (n_phi, n_t1, n_t2)

Return type:

Array

homodyne.optimization.nlsq.fit_computation.compute_g2_batch_with_per_angle_scaling(physical_params, t1, t2, phi_angles, q, L, dt, contrasts, offsets)

Compute g2 with per-angle contrast/offset in single vectorized operation.

Performance Optimization (Spec 006 - FR-007a): Extends compute_g2_batch for per-angle scaling parameters.

Parameters:

• physical_params (Array) – Physical parameters array

• t1 (Array) – Time values

• t2 (Array) – Time values

• phi_angles (Array) – Phi angles in radians, shape (n_phi,)

• q (float) – Experimental parameters

• L (float) – Experimental parameters

• dt (float) – Experimental parameters

• contrasts (Array) – Per-angle contrasts, shape (n_phi,)

• offsets (Array) – Per-angle offsets, shape (n_phi,)

Returns:

g2 values with scaling applied, shape (n_phi, n_t1, n_t2)

Return type:

Array

homodyne.optimization.nlsq.fit_computation.solve_lstsq_batch(theory_batch, exp_batch)

Batch least squares solving for all angles.

Performance Optimization (Spec 006 - FR-008): Vectorized least squares using jax.vmap for all angles simultaneously.

Parameters:

• theory_batch (Array) – Theory values flattened, shape (n_phi, n_t1 * n_t2)

• exp_batch (Array) – Experimental values flattened, shape (n_phi, n_t1 * n_t2)

Returns:

(contrasts, offsets) each shape (n_phi,)

Return type:

tuple[Array, Array]

homodyne.optimization.nlsq.fit_computation.normalize_analysis_mode(mode, n_params, n_angles)

Resolve analysis mode, inferring from parameter counts if needed.

Parameters:

• mode (str | None) – Explicit mode or None

• n_params (int) – Number of parameters

• n_angles (int) – Number of angles

Returns:

‘static’ or ‘laminar_flow’

Return type:

str

homodyne.optimization.nlsq.fit_computation.get_physical_param_count(analysis_mode)

Get number of physical parameters for analysis mode.

Parameters:

analysis_mode (str) – ‘static’ or ‘laminar_flow’

Return type:

int

Returns:

Number of physical parameters

Raises:

ValueError – If mode is unknown

homodyne.optimization.nlsq.fit_computation.extract_parameters_from_result(parameters, n_angles, analysis_mode)

Extract contrast, offset, and physical parameters from result.

Handles both per-angle and scalar parameter layouts.

Parameters:

• parameters (ndarray) – Full parameter array from optimization

• n_angles (int) – Number of phi angles

• analysis_mode (str) – ‘static’ or ‘laminar_flow’

Return type:

tuple[ndarray, ndarray, ndarray, bool]

Returns:

Tuple of (contrasts, offsets, physical_params, scalar_expansion_used)

Raises:

ValueError – If parameter count doesn’t match expected

homodyne.optimization.nlsq.fit_computation.compute_theoretical_fits(result, data, metadata, *, analysis_mode=None, include_solver_surface=True)

Compute theoretical fits with per-angle least squares scaling.

Generates theoretical correlation functions using optimized parameters, then applies per-angle scaling (contrast, offset) via least squares fitting to match experimental intensities.

Parameters:

• result (Any) – NLSQ optimization result with physical parameters

• data (dict[str, Any]) – Experimental data with phi_angles_list, c2_exp, t1, t2

• metadata (dict[str, Any]) – Metadata with L, dt, q for theoretical computation

• analysis_mode (str | None) – Optional analysis mode override

• include_solver_surface (bool) – Whether to include solver surface in output

Returns:

• ‘c2_theoretical_raw’: Raw theoretical fits (n_angles, n_t1, n_t2)

• ’c2_theoretical_scaled’: Scaled fits (n_angles, n_t1, n_t2)

• ’c2_solver_scaled’: Solver surface (if requested)

• ’per_angle_scaling’: Post-hoc lstsq scaling params (n_angles, 2)

• ’per_angle_scaling_solver’: Original solver scaling params

• ’residuals’: Exp - scaled fit (n_angles, n_t1, n_t2)

• ’scalar_per_angle_expansion’: Whether scalar expansion was used

Return type:

dict[str, Any]

Raises:

ValueError – If q is missing or parameter count is invalid

Result Builder

Result building and quality metrics for NLSQ optimization.

Result Building Utilities for NLSQ Optimization.

This module provides utilities for building and processing optimization results, extracted from wrapper.py to improve code organization.

Extracted from wrapper.py as part of refactoring (Dec 2025).

class homodyne.optimization.nlsq.result_builder.QualityMetrics

Bases: object

Quality metrics for optimization results.

chi_squared

Sum of squared residuals

reduced_chi_squared

chi_squared / degrees of freedom

quality_flag

‘good’, ‘marginal’, or ‘poor’

n_at_bounds

Number of parameters at bounds

chi_squared: float

reduced_chi_squared: float

quality_flag: str

n_at_bounds: int = 0

__init__(chi_squared, reduced_chi_squared, quality_flag, n_at_bounds=0)

homodyne.optimization.nlsq.result_builder.compute_quality_metrics(residuals, n_data, n_params, parameter_status=None)

Compute quality metrics from residuals.

Parameters:

• residuals (ndarray) – Array of residuals

• n_data (int) – Number of data points

• n_params (int) – Number of parameters

• parameter_status (list[str] | None) – List of parameter statuses (optional)

Return type:

QualityMetrics

Returns:

QualityMetrics with computed values

homodyne.optimization.nlsq.result_builder.compute_uncertainties(covariance)

Extract parameter uncertainties from covariance matrix.

Parameters:

covariance (ndarray) – Covariance matrix

Return type:

ndarray

Returns:

Array of standard deviations (square root of diagonal)

homodyne.optimization.nlsq.result_builder.normalize_nlsq_result(result, strategy_name='unknown', logger=None)

Normalize various NLSQ result formats to standard format.

NLSQ can return results in different formats depending on the function and version used. This normalizes them to (popt, pcov, info).

Parameters:

• result (Any) – NLSQ result in any format

• strategy_name (str) – Name of strategy for logging

• logger (Any) – Optional logger

Return type:

tuple[ndarray, ndarray, dict[str, Any]]

Returns:

Tuple of (popt, pcov, info)

Raises:

TypeError – If result format is unrecognized

homodyne.optimization.nlsq.result_builder.determine_convergence_status(info, quality_metrics)

Determine convergence status from optimization info.

Parameters:

• info (dict[str, Any]) – Optimization info dict

• quality_metrics (QualityMetrics) – Quality metrics

Returns:

‘converged’, ‘max_iter’, or ‘failed’

Return type:

str

class homodyne.optimization.nlsq.result_builder.ResultBuilder

Bases: object

Builder for constructing OptimizationResult objects.

Provides a fluent interface for building results with proper validation.

parameters: ndarray | None = None

covariance: ndarray | None = None

n_data: int = 0

start_time: float

recovery_actions: list[str]

info: dict[str, Any]

stratification_diagnostics: Any = None

nlsq_diagnostics: dict[str, Any] | None = None

with_parameters(params)

Set optimized parameters.

Return type:

ResultBuilder

with_covariance(cov)

Set parameter covariance matrix.

Return type:

ResultBuilder

with_data_size(n_data)

Set number of data points.

Return type:

ResultBuilder

with_start_time(start_time)

Set optimization start time.

Return type:

ResultBuilder

with_recovery_actions(actions)

Set recovery actions taken.

Return type:

ResultBuilder

with_info(info)

Set optimization info dict.

Return type:

ResultBuilder

with_stratification_diagnostics(diags)

Set stratification diagnostics.

Return type:

ResultBuilder

with_nlsq_diagnostics(diags)

Set NLSQ solver diagnostics.

Return type:

ResultBuilder

with_fourier_covariance_transform(fourier_reparameterizer, n_phi, n_physical)

Transform covariance from Fourier to per-angle space.

T037-T039: Implements Fourier→per-angle covariance transformation.

The transformation uses the Jacobian of the Fourier→per-angle mapping:

Cov_per_angle = J @ Cov_fourier @ J.T

Physical parameter covariance is preserved (not transformed).

Parameters:

• fourier_reparameterizer (Any) – The Fourier reparameterizer used during optimization.

• n_phi (int) – Number of phi angles.

• n_physical (int) – Number of physical parameters.

Returns:

Self for method chaining.

Return type:

ResultBuilder

Notes

If covariance is None or fourier_reparameterizer is None, this method is a no-op.

build(residual_fn=None, xdata=None)

Build the result dictionary.

Parameters:

• residual_fn (Any) – Residual function for computing chi-squared

• xdata (ndarray | None) – X data for residual computation

Return type:

dict[str, Any]

Returns:

Dictionary with all result fields

Raises:

ValueError – If required fields are missing

__init__(parameters=None, covariance=None, n_data=0, start_time=<factory>, recovery_actions=<factory>, info=<factory>, stratification_diagnostics=None, nlsq_diagnostics=None)

Key Classes

homodyne.optimization.nlsq.result_builder.QualityMetrics

Quality metrics for optimization results.

Optimization Strategies

The NLSQ module implements multiple optimization strategies for different dataset sizes:

NLSQ Optimization Strategies Subpackage.

This subpackage contains strategy implementations for NLSQ optimization: - chunking.py: Angle-stratified chunking for large datasets - residual.py: Stratified residual function for per-angle optimization - residual_jit.py: JIT-compiled version of stratified residual - sequential.py: Sequential per-angle optimization - executors.py: Strategy pattern executors for optimization algorithms

NOTE: selection.py (DatasetSizeStrategy, OptimizationStrategy, estimate_memory_requirements) removed in v2.12.0. Use NLSQ’s WorkflowSelector instead.

Chunking Strategy

Angle-Stratified Chunking for Per-Angle Parameter Optimization.

This module implements angle-stratified data reorganization to ensure NLSQ’s chunking strategy remains compatible with per-angle parameters (contrast[i], offset[i] for each phi angle).

Root Cause of Incompatibility:

NLSQ’s chunking splits data arbitrarily without angle awareness. When per-angle parameters are used: - Each contrast[i] only affects points with phi=angle[i] - If a chunk has no points with angle[i], gradient w.r.t. contrast[i] is ZERO - Zero gradients → NLSQ fails silently (0 iterations, unchanged parameters)

Solution: Angle-Stratified Chunking

Reorganize data BEFORE NLSQ optimization so every chunk contains ALL phi angles: - Original: Random 100k-point chunks may miss angles - Stratified: Each 100k-point chunk has balanced angle representation - Result: All per-angle gradients always well-defined

Performance Impact: <1% overhead (0.15s for 3M points) Memory Impact: 2x peak during reorganization (temporary)

Examples

>>> # Reorganize 3M point dataset with 3 angles
>>> phi, t1, t2, g2 = load_data()  # 3M points
>>> phi_s, t1_s, t2_s, g2_s = create_angle_stratified_data(
...     phi, t1, t2, g2, target_chunk_size=100_000
... )
>>> # Now NLSQ optimization will work correctly with per_angle_scaling=True

References

Ultra-Think Analysis: ultra-think-20251106-012247 Issue: Per-angle scaling + NLSQ chunking incompatibility

class homodyne.optimization.nlsq.strategies.chunking.AngleDistributionStats

Bases: object

Statistics about phi angle distribution in dataset.

unique_angles

Array of unique phi angles in the dataset

Type:

np.ndarray

n_angles

Number of unique angles

Type:

int

counts

Points per angle: {angle: count}

Type:

dict[float, int]

fractions

Fraction of total per angle: {angle: fraction}

Type:

dict[float, float]

imbalance_ratio

max(counts) / min(counts), indicates balance

Type:

float

min_angle

Angle with fewest points

Type:

float

max_angle

Angle with most points

Type:

float

is_balanced

True if imbalance_ratio < 5.0 (recommended threshold)

Type:

bool

unique_angles: ndarray

n_angles: int

counts: dict[float, int]

fractions: dict[float, float]

imbalance_ratio: float

min_angle: float

max_angle: float

is_balanced: bool

__init__(unique_angles, n_angles, counts, fractions, imbalance_ratio, min_angle, max_angle, is_balanced)

class homodyne.optimization.nlsq.strategies.chunking.StratificationDiagnostics

Bases: object

Detailed diagnostics for stratification performance and quality.

This dataclass provides comprehensive metrics for analyzing stratification effectiveness, performance, and memory usage.

n_chunks

Number of chunks created

Type:

int

chunk_sizes

Size of each chunk in points

Type:

list[int]

chunk_balance

Chunk size statistics: {mean, std, min, max, cv}

Type:

dict[str, float]

angles_per_chunk

Number of unique angles in each chunk

Type:

list[int]

angle_coverage

Angle coverage statistics: {mean, std, min_coverage_ratio}

Type:

dict[str, float]

execution_time_ms

Time taken for stratification (milliseconds)

Type:

float

memory_overhead_mb

Peak memory overhead during stratification

Type:

float

memory_efficiency

Ratio of data size to peak memory (1.0 = perfect)

Type:

float

throughput_points_per_sec

Processing throughput (points per second)

Type:

float

use_index_based

Whether index-based stratification was used

Type:

bool

n_chunks: int

chunk_sizes: list[int]

chunk_balance: dict[str, float]

angles_per_chunk: list[int]

angle_coverage: dict[str, float]

execution_time_ms: float

memory_overhead_mb: float

memory_efficiency: float

throughput_points_per_sec: float

use_index_based: bool

__init__(n_chunks, chunk_sizes, chunk_balance, angles_per_chunk, angle_coverage, execution_time_ms, memory_overhead_mb, memory_efficiency, throughput_points_per_sec, use_index_based)

homodyne.optimization.nlsq.strategies.chunking.analyze_angle_distribution(phi)

Analyze phi angle distribution to assess balance.

Computes statistics about how data points are distributed across phi angles. This is critical for deciding whether angle-stratified chunking or sequential per-angle optimization should be used.

Parameters:

phi (Array | ndarray) – Array of phi angles (radians or degrees), shape (n_points,)

Returns:

Complete statistics about angle distribution

Return type:

AngleDistributionStats

Examples

>>> phi = np.array([0, 0, 45, 45, 90])  # 2 @ 0°, 2 @ 45°, 1 @ 90°
>>> stats = analyze_angle_distribution(phi)
>>> print(f"Imbalance ratio: {stats.imbalance_ratio:.1f}")
Imbalance ratio: 2.0
>>> print(f"Balanced: {stats.is_balanced}")
Balanced: True

Notes

Imbalance ratio interpretation: - < 2.0: Excellent balance (ideal for stratification) - 2.0 - 5.0: Acceptable balance (stratification works) - > 5.0: High imbalance (consider sequential per-angle) - > 10.0: Very high imbalance (sequential per-angle recommended)

homodyne.optimization.nlsq.strategies.chunking.estimate_stratification_memory(n_points, n_features=4, use_index_based=False, estimated_expansion=1.0)

Estimate memory requirements for stratification ONLY.

WARNING: This function ONLY estimates data reorganization memory. For complete NLSQ optimization memory including Jacobian and optimizer state, use estimate_nlsq_optimization_memory() instead.

Parameters:

• n_points (int) – Total number of data points

• n_features (int) – Number of data features (phi, t1, t2, g2_exp), default: 4

• use_index_based (bool) – If True, use index-based stratification (zero-copy), default: False

• estimated_expansion (float) – Estimated data expansion factor due to Cyclic Stratification (default: 1.0). For imbalanced data, this can be > 1.0 (e.g., 2.0 for 2:1 imbalance).

Returns:

Memory statistics with keys: - original_memory_mb: Original data memory usage - stratified_memory_mb: Memory for stratified copy (including expansion) - peak_memory_mb: Peak memory during stratification - index_memory_mb: Memory for index arrays (if use_index_based) - is_safe: Whether memory usage is safe (<70% of available)

Return type:

dict[str, Any]

Notes

Memory usage: - Full copy: original + (original * expansion) (peak) - Index-based: original + index_array (peak)

homodyne.optimization.nlsq.strategies.chunking.estimate_nlsq_optimization_memory(n_points, n_params, n_features=4, dtype_bytes=8)

Estimate complete memory requirements for NLSQ optimization.

This function provides a COMPLETE memory estimate including all components: - Data arrays (phi, t1, t2, g2) - Jacobian matrix (DOMINANT memory consumer) - JAX JIT compilation overhead - Optimizer internal state

Root Cause Fix (Nov 10, 2025): The original estimate_stratification_memory() only counted data (703 MB), but actual usage was 51 GB (36× underestimate). This function includes ALL memory components for accurate prediction.

Parameters:

• n_points (int) – Total number of data points

• n_params (int) – Number of optimization parameters (e.g., 53 for laminar_flow with per-angle)

• n_features (int) – Number of data features (phi, t1, t2, g2_exp), default: 4

• dtype_bytes (int) – Bytes per floating point number, default: 8 (float64)

Returns:

Complete memory statistics with keys: - data_mb: Data arrays memory - jacobian_mb: Jacobian matrix memory (DOMINANT) - jax_overhead_mb: JAX JIT cache and device arrays - optimizer_mb: Optimizer state (Hessian, gradients) - total_mb: Total estimated memory - peak_gb: Peak memory in GB - available_gb: Available system memory - utilization_pct: Percentage of available memory used - is_safe: Whether memory usage is safe (<70% of available)

Return type:

dict[str, Any]

Examples

>>> # Real dataset from log: 23M points, 53 params
>>> mem = estimate_nlsq_optimization_memory(
...     n_points=23_046_023,
...     n_params=53
... )
>>> print(f"Jacobian: {mem['jacobian_mb']:.0f} MB")
Jacobian: 9,784 MB
>>> print(f"Total: {mem['peak_gb']:.1f} GB")
Total: 14.3 GB
>>> print(f"Utilization: {mem['utilization_pct']:.1f}%")
Utilization: 22.8%
>>>
>>> # With old fixed 100K chunks: 51 GB actual vs 14.3 GB estimated
>>> # Difference due to memory leak (fixed separately)

Notes

Memory Components: 1. Data arrays: n_points × n_features × dtype_bytes 2. Jacobian: n_points × n_params × dtype_bytes (DOMINANT) 3. JAX overhead: 1.75× data (JIT cache, device arrays) 4. Optimizer state: Hessian (n_params²) + gradients + trust region 5. Safety margin: 20% buffer for temporary allocations

Root Cause (Nov 10, 2025): - Old estimate: Only data = 703 MB - Actual peak: 51 GB (includes Jacobian + leak) - New estimate: 14.3 GB (without leak) - With fixes: Expected ~15 GB actual

homodyne.optimization.nlsq.strategies.chunking.calculate_adaptive_chunk_size(total_points, n_params, n_angles, available_memory_gb=None, safety_factor=5.0, min_chunk_size=10_000, max_chunk_size=500_000)

Calculate optimal chunk size based on available system memory and parameter count.

This function addresses the root cause of memory pressure in NLSQ optimization: the fixed 100K chunk size doesn’t account for available memory or the number of parameters, which determines Jacobian matrix size.

The Jacobian matrix dominates memory usage: - Size: n_residuals × n_params × 8 bytes - For 100K points with 53 params: ~42 MB per chunk - Full dataset (23M points): ~9.8 GB Jacobian

Memory Budget Calculation: 1. Reserve 30% for OS, JAX overhead, optimizer state 2. Calculate max points that fit: available_memory / (param_bytes × safety_factor) 3. Ensure all angles fit in each chunk (critical for per-angle parameters) 4. Clamp to reasonable bounds for numerical stability and iteration speed

Parameters:

• total_points (int) – Total number of data points in dataset

• n_params (int) – Number of optimization parameters (e.g., 53 for laminar_flow with per-angle scaling)

• n_angles (int) – Number of unique phi angles (must all fit in each chunk)

• available_memory_gb (float | None) – Available system memory in GB. If None, auto-detected using psutil.

• safety_factor (float) – Multiplicative safety factor for memory overhead (default: 5.0) Accounts for JAX JIT cache, optimizer state, temporary arrays.

• min_chunk_size (int) – Minimum chunk size for numerical stability (default: 10,000)

• max_chunk_size (int) – Maximum chunk size for iteration speed (default: 500,000)

Returns:

Optimal chunk size that fits in available memory

Return type:

int

Examples

>>> # 23M points, 53 parameters, 23 angles, 62GB system
>>> chunk_size = calculate_adaptive_chunk_size(
...     total_points=23_046_023,
...     n_params=53,
...     n_angles=23,
...     available_memory_gb=62.8
... )
>>> print(f"Optimal chunk size: {chunk_size:,}")
Optimal chunk size: 23,000
>>>
>>> # Small dataset, few parameters
>>> chunk_size = calculate_adaptive_chunk_size(
...     total_points=1_000_000,
...     n_params=9,
...     n_angles=3,
...     available_memory_gb=32.0
... )
>>> print(f"Optimal chunk size: {chunk_size:,}")
Optimal chunk size: 500,000  # Clamped to max

Notes

Root Cause Analysis (Nov 10, 2025): - Fixed 100K chunk size caused 96% memory pressure on 62.8GB system - With 53 params: Jacobian alone is 9.8 GB - JAX overhead adds 1.5-2× data size - Optimizer state adds ~2 GB - Total: ~51 GB peak (should be ~15 GB with adaptive sizing)

Algorithm: 1. Auto-detect available memory if not provided 2. Calculate memory per point: n_params × 8 bytes (Jacobian row) 3. Usable memory: 70% of available (reserve 30% for OS/JAX) 4. Max points: usable_memory / (memory_per_point × safety_factor) 5. Chunk size: (max_points / n_angles) × n_angles # Ensure all angles fit 6. Clamp to [min_chunk_size, max_chunk_size]

homodyne.optimization.nlsq.strategies.chunking.create_angle_stratified_data(phi, t1, t2, g2_exp, target_chunk_size=100_000)

Ensure each chunk contains every phi angle using Cyclic Stratification.

Reorders data so NLSQ chunking keeps balanced angle coverage and maintains valid gradients for per-angle parameters.

CRITICAL FIX (Jan 2026): Cyclic Stratification

Previously, stratification stopped when the smallest angle was exhausted, dumping all remaining data into a single massive, unbalanced chunk. This caused rank-deficient Jacobians (zero gradients for missing angles) and memory spikes.

New Logic: 1. Determine number of chunks based on failure mode: ensuring ALL data is used regardless of balance. 2. Iterate through chunks, pulling data from EACH angle. 3. If an angle runs out of data, recycled data from the beginning (Cyclic). 4. Result: Consistent chunk sizes, all angles present in all chunks.

param phi:

Phi angles (radians or degrees), shape (n_points,)

type phi:

Array

param t1:

First time delays, shape (n_points,)

type t1:

Array

param t2:

Second time delays, shape (n_points,)

type t2:

Array

param g2_exp:

Experimental g2 values, shape (n_points,)

type g2_exp:

Array

param target_chunk_size:

Target size for each chunk (default: 100,000) NLSQ typically uses 100k chunks for LARGE/CHUNKED strategies

type target_chunk_size:

int

rtype:

tuple[Array, Array, Array, Array, list[int]]

returns:

• phi_stratified (jnp.ndarray) – Stratified phi angles

• t1_stratified (jnp.ndarray) – Stratified t1 delays

• t2_stratified (jnp.ndarray) – Stratified t2 delays

• g2_stratified (jnp.ndarray) – Stratified g2 values

• chunk_sizes (list[int]) – Size of each stratified chunk (CRITICAL for correct re-chunking)

homodyne.optimization.nlsq.strategies.chunking.create_angle_stratified_indices(phi, target_chunk_size=100_000)

Create index array for zero-copy angle-stratified data access using Interleaved Stratification.

This function implements index-based stratification, reducing memory overhead from 2x (full copy) to ~1% (index array only).

Interleaved Stratification

Distributes data from each angle group across chunks using round-robin allocation. Each angle’s data is split proportionally across chunks, ensuring: - No data expansion (output size = input size) - No duplicates - All angles represented in each chunk (for balanced data)

param phi:

Phi angles (radians or degrees), shape (n_points,)

type phi:

Array | ndarray

param target_chunk_size:

Target size for each chunk (default: 100,000)

type target_chunk_size:

int

rtype:

tuple[ndarray, list[int]]

returns:

• indices (np.ndarray) – Index array specifying stratified ordering, shape (n_points,) Use: data_stratified = data_original[indices]

• chunk_sizes (list[int]) – Size of each stratified chunk (CRITICAL for correct re-chunking)

class homodyne.optimization.nlsq.strategies.chunking.StratifiedIndexIterator

Bases: object

Iterator that yields index chunks for stratified data access.

This iterator allows processing strictly stratified chunks one by one without materializing the full index array or data chunks in memory.

indices: ndarray

chunk_sizes: list[int]

__init__(indices, chunk_sizes)

homodyne.optimization.nlsq.strategies.chunking.get_stratified_chunk_iterator(phi, target_chunk_size=100_000)

Create an iterator yielding stratified index chunks.

Parameters:

• phi (Array | ndarray) – Array of phi angles

• target_chunk_size (int) – Desired chunk size

Return type:

StratifiedIndexIterator

Returns:

StratifiedIndexIterator yielding index chunks

homodyne.optimization.nlsq.strategies.chunking.should_use_stratification(n_points, n_angles, per_angle_scaling, imbalance_ratio)

Decide whether to use angle-stratified chunking.

Decision logic: - Small datasets (<100k): No (use STANDARD strategy, no chunking) - No per-angle scaling: No (regular chunking works fine) - High imbalance (>5:1): No (use sequential per-angle instead) - Otherwise: Yes (use stratified chunking)

Parameters:

• n_points (int) – Total number of data points

• n_angles (int) – Number of unique phi angles

• per_angle_scaling (bool) – Whether per-angle parameters are enabled

• imbalance_ratio (float) – max(angle_counts) / min(angle_counts)

Return type:

tuple[bool, str]

Returns:

• should_stratify (bool) – True if stratification should be used

• reason (str) – Human-readable explanation of decision

Examples

>>> should, reason = should_use_stratification(
...     n_points=3_000_000,
...     n_angles=3,
...     per_angle_scaling=True,
...     imbalance_ratio=2.5
... )
>>> print(should, reason)
True "Large dataset with balanced angles"

homodyne.optimization.nlsq.strategies.chunking.compute_stratification_diagnostics(phi_original, phi_stratified, execution_time_ms, use_index_based=False, target_chunk_size=100_000, chunk_sizes=None)

Compute detailed diagnostics for stratification quality and performance.

This function analyzes the stratified data to provide comprehensive metrics about chunk balance, angle coverage, memory efficiency, and throughput.

Parameters:

• phi_original (ndarray) – Original phi angles before stratification

• phi_stratified (ndarray) – Stratified phi angles after reorganization

• execution_time_ms (float) – Time taken for stratification (milliseconds)

• use_index_based (bool) – Whether index-based stratification was used, default: False

• target_chunk_size (int) – Target chunk size used, default: 100,000

Returns:

Comprehensive diagnostic metrics

Return type:

StratificationDiagnostics

Examples

>>> import time
>>> phi = np.repeat([0, 45, 90], 100)
>>> start = time.perf_counter()
>>> phi_s, t1_s, t2_s, g2_s = create_angle_stratified_data(phi, t1, t2, g2)
>>> exec_time_ms = (time.perf_counter() - start) * 1000
>>> diagnostics = compute_stratification_diagnostics(
...     phi, phi_s, exec_time_ms, use_index_based=False
... )
>>> print(f"Chunks: {diagnostics.n_chunks}")
>>> print(f"Throughput: {diagnostics.throughput_points_per_sec:,.0f} pts/s")

homodyne.optimization.nlsq.strategies.chunking.format_diagnostics_report(diagnostics)

Format stratification diagnostics as human-readable report.

Parameters:

diagnostics (StratificationDiagnostics) – Diagnostic metrics to format

Returns:

Formatted report with all diagnostic metrics

Return type:

str

Examples

>>> diagnostics = compute_stratification_diagnostics(phi, phi_s, 150.0)
>>> report = format_diagnostics_report(diagnostics)
>>> print(report)

Residual Functions

class homodyne.optimization.nlsq.strategies.residual.StratifiedResidualFunction

Bases: object

Residual function that respects angle-stratified chunk structure.

This class wraps the model’s residual computation to work with stratified chunks, ensuring that each chunk contains all phi angles. This is critical for per-angle scaling parameters to have non-zero gradients.

The function is designed to work with NLSQ’s least_squares() function, which calls the residual function at each optimization iteration.

chunks

List of angle-stratified data chunks

model

TheoryEngine instance for computing residuals

per_angle_scaling

Whether per-angle scaling is enabled

logger

Logger instance for diagnostics

n_chunks

Number of stratified chunks

n_total_points

Total number of data points across all chunks

compute_chunk_jit

JIT-compiled chunk residual computation

__init__(stratified_data, per_angle_scaling, physical_param_names, logger=None)

Initialize the stratified residual function.

Parameters:

• stratified_data (Any) – Object with .chunks attribute containing angle-stratified chunks. Each chunk must have: phi, t1, t2, g2, q, L, dt attributes. stratified_data.sigma contains the full 3D sigma array (metadata).

• per_angle_scaling (bool) – Whether per-angle scaling parameters are used.

• physical_param_names (list[str]) – List of physical parameter names (e.g., [‘D0’, ‘alpha’, ‘D_offset’])

• logger (Logger | None) – Optional logger for diagnostics.

Raises:

ValueError – If stratified_data.chunks is empty or invalid.

jax_residual(params)

Return type:

Array

validate_chunk_structure()

Validate that all chunks contain all phi angles.

This is a critical validation to ensure per-angle parameter gradients will be non-zero. If any chunk is missing an angle, the gradient for that angle’s parameters will be zero, causing optimization failure.

Return type:

bool

Returns:

True if validation passes

Raises:

ValueError – If any chunk is missing angles or has inconsistent structure

get_diagnostics()

Get diagnostic information about the residual function.

Returns:

• n_chunks: Number of chunks

• n_total_points: Total data points

• n_angles: Number of unique phi angles

• per_angle_scaling: Whether per-angle scaling is enabled

• chunk_sizes: List of points per chunk

• chunk_angle_counts: List of angles per chunk

• min_chunk_size: Minimum chunk size

• max_chunk_size: Maximum chunk size

• mean_chunk_size: Mean chunk size

Return type:

dict[str, Any]

log_diagnostics()

Log diagnostic information for monitoring.

Return type:

None

homodyne.optimization.nlsq.strategies.residual.create_stratified_residual_function(stratified_data, per_angle_scaling, physical_param_names, logger=None, validate=True)

Factory function to create and validate a stratified residual function.

This is a convenience function that creates a StratifiedResidualFunction, optionally validates its structure, and logs diagnostics.

Parameters:

• stratified_data (Any) – Object with .chunks attribute containing angle-stratified chunks

• per_angle_scaling (bool) – Whether per-angle scaling parameters are used

• physical_param_names (list[str]) – List of physical parameter names (e.g., [‘D0’, ‘alpha’, ‘D_offset’])

• logger (Logger | None) – Optional logger for diagnostics

• validate (bool) – Whether to validate chunk structure (recommended)

Return type:

StratifiedResidualFunction

Returns:

Validated StratifiedResidualFunction instance

Raises:

ValueError – If validation fails

Example

>>> residual_fn = create_stratified_residual_function(
...     stratified_data=stratified_data,
...     per_angle_scaling=True,
...     physical_param_names=['D0', 'alpha', 'D_offset'],
...     validate=True
... )
>>> residual_fn.log_diagnostics()

JIT-Compiled Residual Functions

JIT-compatible stratified residual function using padded vmap for full JIT compilation.

JAX JIT-compatible stratified residual function for NLSQ optimization.

This module provides a JIT-compatible version of StratifiedResidualFunction that uses static shapes and vmap for vectorization, solving the JAX tracing incompatibility.

Key Improvements over original StratifiedResidualFunction: - Uses jax.vmap for parallel chunk processing (no Python loops) - Pads chunks to uniform size for static shapes (JIT-compatible) - Fully JIT-compiled for maximum performance - Maintains angle stratification guarantee

Author: Homodyne Development Team Date: 2025-11-13 Version: 2.4.0

class homodyne.optimization.nlsq.strategies.residual_jit.StratifiedResidualFunctionJIT

Bases: object

JIT-compatible stratified residual function using padded vmap.

This class solves the JAX JIT incompatibility by: 1. Padding all chunks to uniform size (static shapes) 2. Using jax.vmap for vectorized parallel processing 3. Masking padded values in the final residuals

The function maintains angle stratification (all chunks contain all angles) while being fully JIT-compilable.

phi_padded

Padded phi arrays (n_chunks, max_chunk_size)

t1_padded

Padded t1 arrays (n_chunks, max_chunk_size)

t2_padded

Padded t2 arrays (n_chunks, max_chunk_size)

g2_padded

Padded g2 observations (n_chunks, max_chunk_size)

mask

Boolean mask for real vs padded data (n_chunks, max_chunk_size)

n_chunks

Number of stratified chunks

max_chunk_size

Maximum points per chunk (for padding)

n_real_points

Total number of real (non-padded) data points

__init__(stratified_data, per_angle_scaling, physical_param_names, logger=None, fixed_contrast_per_angle=None, fixed_offset_per_angle=None)

Initialize JIT-compatible stratified residual function.

Parameters:

• stratified_data (Any) – Object with .chunks attribute containing angle-stratified chunks

• per_angle_scaling (bool) – Whether per-angle scaling parameters are used

• physical_param_names (list[str]) – List of physical parameter names

• logger (Logger | None) – Optional logger for diagnostics

• fixed_contrast_per_angle (ndarray | None) – Fixed per-angle contrast values (for constant mode). When provided, contrast is NOT included in the parameter vector.

• fixed_offset_per_angle (ndarray | None) – Fixed per-angle offset values (for constant mode). When provided, offset is NOT included in the parameter vector.

__call__(params)

Compute residuals (interface for NLSQ least_squares).

This method is JIT-traced by NLSQ, so it must use JAX operations only. Padded values are already masked to zero, so they don’t contribute to the optimization objective (sum of squared residuals).

Parameters:

params (ndarray | Array) – Parameters (numpy or JAX array)

Return type:

Array

Returns:

Residuals as JAX array (n_chunks * max_chunk_size,) with zeros for padding Note: Padding zeros don’t affect optimization but increase array size

validate_chunk_structure()

Validate that all chunks contain all phi angles.

Return type:

bool

Returns:

True if validation passes

Raises:

ValueError – If validation fails

get_diagnostics()

Get diagnostic information about the residual function.

Return type:

dict

log_diagnostics()

Log diagnostic information about the residual function.

Return type:

None

Key Classes

homodyne.optimization.nlsq.strategies.residual_jit.StratifiedResidualFunctionJIT

JIT-compatible stratified residual function using padded vmap.

Key Features

• Static shapes: Pads chunks to uniform size for JIT compatibility

• vmap vectorization: Parallel chunk processing without Python loops

• Angle stratification: Maintains all angles in each chunk

Sequential Optimization

Sequential Per-Angle Optimization Module

Provides fallback optimization strategy when angle-stratified chunking cannot be used. Optimizes each phi angle independently and combines results.

Use Cases: - Extreme angle imbalance (ratio > 5.0) - Stratification explicitly disabled - Debugging and validation - Memory-constrained environments

Author: Homodyne Development Team Version: 2.3.0 Date: 2026-01-14

class homodyne.optimization.nlsq.strategies.sequential.AngleSubset

Bases: object

Data subset for a single phi angle.

phi_angle

The phi angle value for this subset

Type:

float

phi_indices

Indices where phi == phi_angle

Type:

np.ndarray

n_points

Number of data points for this angle

Type:

int

phi

Phi values (all equal to phi_angle)

Type:

np.ndarray

t1

Time 1 values

Type:

np.ndarray

t2

Time 2 values

Type:

np.ndarray

g2_exp

Experimental g2 values

Type:

np.ndarray

phi_angle: float

phi_indices: ndarray

n_points: int

phi: ndarray

t1: ndarray

t2: ndarray

g2_exp: ndarray

__init__(phi_angle, phi_indices, n_points, phi, t1, t2, g2_exp)

class homodyne.optimization.nlsq.strategies.sequential.SequentialResult

Bases: object

Result from sequential per-angle optimization.

combined_parameters

Combined optimized parameters (weighted average)

Type:

np.ndarray

combined_covariance

Combined covariance matrix

Type:

np.ndarray

per_angle_results

Individual results for each angle

Type:

list[dict]

n_angles_optimized

Number of angles successfully optimized

Type:

int

n_angles_failed

Number of angles that failed optimization

Type:

int

total_cost

Combined optimization cost

Type:

float

success_rate

Fraction of angles that converged (0.0-1.0)

Type:

float

combined_parameters: ndarray

combined_covariance: ndarray

per_angle_results: list[dict[str, Any]]

n_angles_optimized: int

n_angles_failed: int

total_cost: float

success_rate: float

initial_jacobian_norms: ndarray | None = None

final_jacobian_norms: ndarray | None = None

__init__(combined_parameters, combined_covariance, per_angle_results, n_angles_optimized, n_angles_failed, total_cost, success_rate, initial_jacobian_norms=None, final_jacobian_norms=None)

homodyne.optimization.nlsq.strategies.sequential.split_data_by_angle(phi, t1, t2, g2_exp, min_points_per_angle=10)

Split dataset into per-angle subsets.

Parameters:

• phi (ndarray) – Phi angle values (flattened)

• t1 (ndarray) – Time 1 values (flattened)

• t2 (ndarray) – Time 2 values (flattened)

• g2_exp (ndarray) – Experimental g2 values (flattened)

• min_points_per_angle (int) – Minimum points required per angle, default: 10

Returns:

List of angle subsets, one per unique phi value

Return type:

list[AngleSubset]

Raises:

ValueError – If any angle has fewer than min_points_per_angle points

Examples

>>> phi = np.array([0, 0, 90, 90, 180, 180])
>>> t1 = np.linspace(0, 1, 6)
>>> t2 = np.linspace(0, 1, 6)
>>> g2 = np.ones(6)
>>> subsets = split_data_by_angle(phi, t1, t2, g2)
>>> len(subsets)
3
>>> subsets[0].phi_angle
0.0
>>> subsets[0].n_points
2

homodyne.optimization.nlsq.strategies.sequential.optimize_single_angle(subset, residual_func, initial_params, bounds, **optimizer_kwargs)

Optimize parameters for a single phi angle.

Parameters:

• subset (AngleSubset) – Data for this angle

• residual_func (Callable) – Residual function: residual_func(params, phi, t1, t2) -> residuals

• initial_params (ndarray) – Initial parameter guess

• bounds (tuple[ndarray, ndarray]) – (lower_bounds, upper_bounds) for parameters

• **optimizer_kwargs – Additional arguments passed to NLSQ optimizer

Returns:

Result dictionary with keys: - ‘parameters’: Optimized parameters - ‘covariance’: Covariance matrix - ‘cost’: Final cost - ‘success’: Whether optimization converged - ‘n_iterations’: Number of iterations - ‘message’: Status message - ‘n_points’: Number of points used - ‘phi_angle’: Angle value

Return type:

dict[str, Any]

Notes

Uses NLSQ LeastSquares for JAX-accelerated optimization.

homodyne.optimization.nlsq.strategies.sequential.combine_angle_results(per_angle_results, weighting='inverse_variance')

Combine per-angle optimization results.

Parameters:

• per_angle_results (list[dict[str, Any]]) – Results from optimize_single_angle for each angle

• weighting (str) – Weighting scheme: ‘inverse_variance’ | ‘uniform’ | ‘n_points’ Default: ‘inverse_variance’ (optimal statistical weighting)

Return type:

tuple[ndarray, ndarray, float]

Returns:

• combined_params (np.ndarray) – Weighted average of parameters

• combined_cov (np.ndarray) – Combined covariance matrix

• total_cost (float) – Sum of individual costs

Notes

Inverse variance weighting:

w_i = 1 / σ²_i μ = Σ(w_i × x_i) / Σ(w_i) σ² = 1 / Σ(w_i)

This provides optimal statistical combination when errors are independent.

homodyne.optimization.nlsq.strategies.sequential.strip_fixed_parameters(initial_params, lower_bounds, upper_bounds)

Remove fixed parameters (lower == upper) from the optimizer inputs.

The TRF solver used by sequential optimization requires strict lower < upper for every parameter. Fixed parameters (equality constraints encoded as lower == upper) must be stripped before the call and their known values re-inserted into the result.

Parameters:

• initial_params (ndarray) – Full parameter vector including fixed parameters.

• lower_bounds (ndarray) – Lower bounds array (same length as initial_params).

• upper_bounds (ndarray) – Upper bounds array (same length as initial_params).

Return type:

tuple[ndarray, ndarray, ndarray, ndarray]

Returns:

• free_params (np.ndarray) – Subset of initial_params where lower < upper.

• free_lower (np.ndarray) – Lower bounds for free parameters.

• free_upper (np.ndarray) – Upper bounds for free parameters.

• free_mask (np.ndarray) – Boolean mask (length == len(initial_params)), True where free.

Examples

>>> p = np.array([1.0, 2.0, 3.0])
>>> lo = np.array([0.0, 2.0, 0.0])
>>> hi = np.array([5.0, 2.0, 5.0])
>>> free, fl, fu, mask = strip_fixed_parameters(p, lo, hi)
>>> free       # array([1.0, 3.0])
>>> mask       # array([True, False, True])

homodyne.optimization.nlsq.strategies.sequential.restore_fixed_parameters(free_result, fixed_values, free_mask)

Re-insert fixed parameter values into the optimized result.

Inverse of strip_fixed_parameters().

Parameters:

• free_result (ndarray) – Optimized values for the free parameters.

• fixed_values (ndarray) – Full reference parameter vector (fixed positions taken from here).

• free_mask (ndarray) – Boolean mask returned by strip_fixed_parameters().

Returns:

Full parameter vector with fixed values restored.

Return type:

ndarray

homodyne.optimization.nlsq.strategies.sequential.optimize_per_angle_sequential(phi, t1, t2, g2_exp, residual_func, initial_params, bounds, weighting='inverse_variance', min_success_rate=0.5, parameter_names=None, **optimizer_kwargs)

Optimize parameters sequentially for each phi angle.

Main entry point for sequential per-angle optimization.

Parameters:

• phi (ndarray) – Phi angle values (flattened)

• t1 (ndarray) – Time 1 values (flattened)

• t2 (ndarray) – Time 2 values (flattened)

• g2_exp (ndarray) – Experimental g2 values (flattened)

• residual_func (callable) – Residual function: residual_func(params, phi, t1, t2, g2) -> residuals

• initial_params (ndarray) – Initial parameter guess

• bounds (tuple[ndarray, ndarray]) – (lower_bounds, upper_bounds)

• weighting (str) – Result combination weighting: ‘inverse_variance’ | ‘uniform’ | ‘n_points’

• min_success_rate (float) – Minimum fraction of angles that must converge (0.0-1.0), default: 0.5

• parameter_names (Sequence[str] | None) – Parameter ordering used to align per-parameter kwargs (e.g., x_scale)

• **optimizer_kwargs – Additional arguments passed to NLSQ LeastSquares.least_squares

Returns:

Combined optimization results

Return type:

SequentialResult

Raises:

RuntimeError – If success rate < min_success_rate

Examples

>>> # Simple example with 3 angles
>>> phi = np.array([0]*100 + [90]*100 + [180]*100)
>>> t1 = np.tile(np.linspace(0, 1, 100), 3)
>>> t2 = np.tile(np.linspace(0, 1, 100), 3)
>>> g2 = np.ones(300)
>>>
>>> def residuals(params, phi, t1, t2, g2):
...     # Simple model
...     return g2 - (1.0 + params[0] * np.exp(-params[1] * t1))
>>>
>>> result = optimize_per_angle_sequential(
...     phi, t1, t2, g2,
...     residuals,
...     initial_params=np.array([0.5, 1.0]),
...     bounds=(np.array([0.0, 0.0]), np.array([1.0, 10.0]))
... )
>>> result.success_rate
1.0
>>> len(result.per_angle_results)
3

Strategy Executors

Implementation of the Strategy pattern for optimization execution.

Optimization Strategy Executors for NLSQ.

This module implements the Strategy pattern for different optimization approaches, enabling cleaner code organization and easier testing.

Extracted from wrapper.py as part of refactoring (Dec 2025).

class homodyne.optimization.nlsq.strategies.executors.ExecutionResult

Bases: object

Result from optimization execution.

popt

Optimized parameters

pcov

Parameter covariance matrix

info

Additional optimization information

recovery_actions

List of recovery actions taken

convergence_status

‘converged’, ‘partial’, or ‘failed’

popt: ndarray

pcov: ndarray

info: dict[str, Any]

recovery_actions: list[str]

convergence_status: str

__init__(popt, pcov, info, recovery_actions, convergence_status)

class homodyne.optimization.nlsq.strategies.executors.OptimizationExecutor

Bases: ABC

Abstract base class for optimization strategy executors.

Implements the Strategy pattern for different optimization approaches. Each concrete implementation handles a specific optimization method.

abstractmethod execute(residual_fn, xdata, ydata, initial_params, bounds, loss_name, x_scale_value, logger)

Execute optimization with the specific strategy.

Parameters:

• residual_fn (Callable[..., Any]) – Residual function to minimize

• xdata (ndarray) – Independent variable data

• ydata (ndarray) – Dependent variable data (observations)

• initial_params (ndarray) – Initial parameter guess

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds as (lower, upper) tuple

• loss_name (str) – Loss function name (e.g., ‘soft_l1’)

• x_scale_value (float | ndarray | str) – Parameter scaling for trust region

• logger (Any) – Logger instance

Return type:

ExecutionResult

Returns:

ExecutionResult with optimized parameters and diagnostics

abstract property name: str

Strategy name for logging.

abstract property supports_progress: bool

Whether this strategy supports progress bars.

class homodyne.optimization.nlsq.strategies.executors.StandardExecutor

Bases: OptimizationExecutor

Standard curve_fit optimization for small datasets (<1M points).

Uses scipy.optimize.curve_fit through the NLSQ wrapper. Fast for small datasets, but doesn’t handle large datasets efficiently.

property name: str

Strategy name for logging.

property supports_progress: bool

Whether this strategy supports progress bars.

execute(residual_fn, xdata, ydata, initial_params, bounds, loss_name, x_scale_value, logger)

Execute standard curve_fit optimization.

Return type:

ExecutionResult

class homodyne.optimization.nlsq.strategies.executors.LargeDatasetExecutor

Bases: OptimizationExecutor

Large dataset optimization using curve_fit_large.

Uses NLSQ’s memory-efficient curve_fit_large function for datasets that exceed memory limits with standard curve_fit.

property name: str

Strategy name for logging.

property supports_progress: bool

Whether this strategy supports progress bars.

execute(residual_fn, xdata, ydata, initial_params, bounds, loss_name, x_scale_value, logger)

Execute large dataset optimization.

Return type:

ExecutionResult

class homodyne.optimization.nlsq.strategies.executors.StreamingExecutor

Bases: OptimizationExecutor

Streaming optimization for unlimited dataset sizes.

Uses NLSQ’s AdaptiveHybridStreamingOptimizer for datasets that are too large to fit in memory. Supports checkpointing and recovery.

Note

The old StreamingOptimizer was removed in NLSQ 0.4.0. This executor now uses AdaptiveHybridStreamingOptimizer which provides better convergence and parameter estimation.

__init__(checkpoint_config=None)

Initialize streaming executor.

Parameters:

checkpoint_config (dict[str, Any] | None) – Configuration for checkpointing and hybrid streaming

property name: str

Strategy name for logging.

property supports_progress: bool

Whether this strategy supports progress bars.

execute(residual_fn, xdata, ydata, initial_params, bounds, loss_name, x_scale_value, logger)

Execute streaming optimization using AdaptiveHybridStreamingOptimizer.

Return type:

ExecutionResult

homodyne.optimization.nlsq.strategies.executors.get_executor(strategy_name, checkpoint_config=None)

Factory function to get the appropriate executor.

Parameters:

• strategy_name (str) – Name of strategy (‘standard’, ‘large’, ‘streaming’)

• checkpoint_config (dict[str, Any] | None) – Configuration for streaming checkpoints

Return type:

OptimizationExecutor

Returns:

OptimizationExecutor instance for the strategy

Raises:

ValueError – If strategy name is unknown

Key Classes

homodyne.optimization.nlsq.strategies.executors.ExecutionResult	Result from optimization execution.
homodyne.optimization.nlsq.strategies.executors.OptimizationExecutor	Abstract base class for optimization strategy executors.
homodyne.optimization.nlsq.strategies.executors.StandardExecutor	Standard curve_fit optimization for small datasets (<1M points).
homodyne.optimization.nlsq.strategies.executors.LargeDatasetExecutor	Large dataset optimization using curve_fit_large.
homodyne.optimization.nlsq.strategies.executors.StreamingExecutor	Streaming optimization for unlimited dataset sizes.

Multi-Start Optimization

Multi-start optimization explores parameter space from multiple starting points using Latin Hypercube Sampling to find the global optimum and detect parameter degeneracy.

Multi-start NLSQ optimization with Latin Hypercube Sampling.

This module implements multi-start optimization to explore the parameter space and avoid local minima. All datasets use the FULL strategy (N complete fits).

NOTE: Subsampling is explicitly NOT supported per project requirements. Numerical precision and reproducibility take priority over computational speed.

Part of homodyne v2.6.0 architecture.

class homodyne.optimization.nlsq.multistart.MultiStartConfig

Bases: object

Configuration for multi-start optimization.

enable

Whether to use multi-start optimization. Default: False.

Type:

bool

n_starts

Number of starting points to generate. Default: 10.

Type:

int

seed

Random seed for reproducibility. Default: 42.

Type:

int

sampling_strategy

Method for generating starting points. “latin_hypercube” or “random”.

Type:

str

custom_starts

User-provided custom starting points to include alongside generated starts.

Type:

list[list[float]] | None

n_workers

Number of parallel workers. 0 = auto (min of n_starts, cpu_count).

Type:

int

use_screening

Whether to pre-filter starting points by initial cost.

Type:

bool

screen_keep_fraction

Fraction of starting points to keep after screening.

Type:

float

refine_top_k

Number of top solutions to refine with tighter tolerance.

Type:

int

refinement_ftol

Function tolerance for refinement phase.

Type:

float

degeneracy_threshold

Chi-squared similarity threshold for degeneracy detection.

Type:

float

enable: bool = False

n_starts: int = 10

seed: int = 42

sampling_strategy: str = 'latin_hypercube'

custom_starts: list[list[float]] | None = None

n_workers: int = 0

use_screening: bool = True

screen_keep_fraction: float = 0.5

refine_top_k: int = 3

refinement_ftol: float = 1e-12

degeneracy_threshold: float = 0.1

classmethod from_nlsq_config(nlsq_config)

Create MultiStartConfig from NLSQConfig.

Parameters:

nlsq_config (Any) – NLSQ configuration object.

Returns:

Multi-start configuration.

Return type:

MultiStartConfig

to_nlsq_global_config()

Convert to NLSQ’s GlobalOptimizationConfig.

Returns:

NLSQ global optimization configuration.

Return type:

Any

Raises:

ImportError – If NLSQ GlobalOptimizationConfig is not available.

Notes

Maps homodyne’s multi-start configuration to NLSQ’s GlobalOptimizationConfig: - sampling_strategy -> sampler (lhs, sobol, halton) - use_screening -> elimination_rounds (0 if disabled) - screen_keep_fraction -> elimination_fraction (inverted)

__init__(enable=False, n_starts=10, seed=42, sampling_strategy='latin_hypercube', custom_starts=None, n_workers=0, use_screening=True, screen_keep_fraction=0.5, refine_top_k=3, refinement_ftol=1e-12, degeneracy_threshold=0.1)

class homodyne.optimization.nlsq.multistart.SingleStartResult

Bases: object

Result from a single starting point optimization.

start_idx

Index of the starting point in the LHS sequence.

Type:

int

initial_params

Initial parameter values used.

Type:

NDArray[np.float64]

final_params

Optimized parameter values.

Type:

NDArray[np.float64]

chi_squared

Final chi-squared value.

Type:

float

reduced_chi_squared

Chi-squared divided by degrees of freedom.

Type:

float

success

Whether optimization converged successfully.

Type:

bool

status

Optimizer status code.

Type:

int

message

Optimizer status message.

Type:

str

n_iterations

Number of optimization iterations.

Type:

int

n_fev

Number of function evaluations.

Type:

int

wall_time

Execution time in seconds.

Type:

float

hessian

Hessian matrix at solution (for CMC initialization).

Type:

NDArray[np.float64] | None

covariance

Parameter covariance matrix.

Type:

NDArray[np.float64] | None

jacobian

Final Jacobian matrix.

Type:

NDArray[np.float64] | None

start_idx: int

initial_params: ndarray[tuple[Any, ...], dtype[float64]]

final_params: ndarray[tuple[Any, ...], dtype[float64]]

chi_squared: float

reduced_chi_squared: float = inf

success: bool = False

status: int = 0

message: str = ''

n_iterations: int = 0

n_fev: int = 0

wall_time: float = 0.0

hessian: ndarray[tuple[Any, ...], dtype[float64]] | None = None

covariance: ndarray[tuple[Any, ...], dtype[float64]] | None = None

jacobian: ndarray[tuple[Any, ...], dtype[float64]] | None = None

__init__(start_idx, initial_params, final_params, chi_squared, reduced_chi_squared=inf, success=False, status=0, message='', n_iterations=0, n_fev=0, wall_time=0.0, hessian=None, covariance=None, jacobian=None)

class homodyne.optimization.nlsq.multistart.MultiStartResult

Bases: object

Aggregated results from multi-start optimization.

best

Best result by chi-squared.

Type:

SingleStartResult

all_results

All optimization results.

Type:

list[SingleStartResult]

config

Configuration used.

Type:

MultiStartConfig

strategy_used

Strategy that was used (always “full”).

Type:

str

n_successful

Number of successful optimizations.

Type:

int

n_unique_basins

Number of distinct local minima found.

Type:

int

degeneracy_detected

Whether parameter degeneracy was detected.

Type:

bool

total_wall_time

Total execution time in seconds.

Type:

float

screening_costs

Initial costs from screening phase.

Type:

NDArray[np.float64] | None

basin_labels

Cluster labels for each result.

Type:

NDArray[np.int64] | None

best: SingleStartResult

all_results: list[SingleStartResult]

config: MultiStartConfig

strategy_used: str

n_successful: int = 0

n_unique_basins: int = 1

degeneracy_detected: bool = False

total_wall_time: float = 0.0

screening_costs: ndarray[tuple[Any, ...], dtype[float64]] | None = None

basin_labels: ndarray[tuple[Any, ...], dtype[int64]] | None = None

to_optimization_result()

Convert MultiStartResult to OptimizationResult for CLI compatibility.

Returns:

Optimization result object containing the best solution with multi-start metadata in nlsq_diagnostics.

Return type:

OptimizationResult

__init__(best, all_results, config, strategy_used, n_successful=0, n_unique_basins=1, degeneracy_detected=False, total_wall_time=0.0, screening_costs=None, basin_labels=None)

homodyne.optimization.nlsq.multistart.check_zero_volume_bounds(bounds)

Check if parameter bounds have zero volume (all lower == upper).

Parameters:

bounds (ndarray[tuple[Any, ...], dtype[double]]) – Parameter bounds as (n_params, 2) array with [lower, upper] for each.

Returns:

True if bounds have zero volume (all parameters fixed).

Return type:

bool

homodyne.optimization.nlsq.multistart.validate_n_starts_for_lhs(n_starts, n_params, warn=True)

Validate n_starts for Latin Hypercube Sampling coverage.

For LHS to provide meaningful coverage, n_starts should be at least n_params. Very large n_starts relative to parameter space may produce redundant samples.

Parameters:

• n_starts (int) – Requested number of starting points.

• n_params (int) – Number of parameters (dimensions).

• warn (bool) – Whether to emit warnings for suboptimal settings.

Returns:

Validated n_starts (unchanged if valid).

Return type:

int

homodyne.optimization.nlsq.multistart.generate_lhs_starts(bounds, n_starts, seed=42)

Generate starting points via Latin Hypercube Sampling.

Parameters:

• bounds (ndarray[tuple[Any, ...], dtype[double]]) – Parameter bounds as (n_params, 2) array with [lower, upper] for each.

• n_starts (int) – Number of starting points to generate.

• seed (int) – Random seed for reproducibility.

Returns:

Starting points as (n_starts, n_params) array.

Return type:

ndarray[tuple[Any, ...], dtype[double]]

homodyne.optimization.nlsq.multistart.generate_random_starts(bounds, n_starts, seed=42)

Generate starting points via random uniform sampling.

Parameters:

• bounds (ndarray[tuple[Any, ...], dtype[double]]) – Parameter bounds as (n_params, 2) array.

• n_starts (int) – Number of starting points to generate.

• seed (int) – Random seed for reproducibility.

Returns:

Starting points as (n_starts, n_params) array.

Return type:

ndarray[tuple[Any, ...], dtype[double]]

homodyne.optimization.nlsq.multistart.include_custom_starts(generated_starts, custom_starts, bounds)

Include user-provided custom starting points alongside generated starts.

Custom starting points are prepended to the generated starts so they are always included (not filtered by screening).

Parameters:

• generated_starts (ndarray[tuple[Any, ...], dtype[double]]) – Starting points generated by LHS or random sampling.

• custom_starts (list[list[float]] | ndarray[tuple[Any, ...], dtype[double]] | None) – User-provided custom starting points.

• bounds (ndarray[tuple[Any, ...], dtype[double]]) – Parameter bounds for validation.

Returns:

Combined starting points with custom starts first.

Return type:

ndarray[tuple[Any, ...], dtype[double]]

homodyne.optimization.nlsq.multistart.screen_starts(cost_func, starts, keep_fraction=0.5, min_keep=3, n_workers=0)

Pre-filter starting points by initial cost.

Parameters:

• cost_func (Callable[[ndarray[tuple[Any, ...], dtype[double]]], float]) – Function that computes cost (chi-squared) for a parameter vector.

• starts (ndarray[tuple[Any, ...], dtype[double]]) – Starting points as (n_starts, n_params) array.

• keep_fraction (float) – Fraction of starting points to keep (0, 1].

• min_keep (int) – Minimum number of starting points to keep.

• n_workers (int) – Number of parallel workers for cost evaluation. 0 = auto (cpu_count - 1).

Returns:

Filtered starting points and their initial costs.

Return type:

tuple[ndarray[tuple[Any, ...], dtype[double]], ndarray[tuple[Any, ...], dtype[double]]]

homodyne.optimization.nlsq.multistart.detect_degeneracy(results, chi_sq_threshold=0.1, param_threshold=0.2)

Detect parameter degeneracy from multiple optimization results.

Parameters:

• results (list[SingleStartResult]) – List of optimization results.

• chi_sq_threshold (float) – Maximum relative chi-squared difference to consider similar.

• param_threshold (float) – Maximum relative parameter distance to consider same basin.

Returns:

(degeneracy_detected, n_unique_basins, basin_labels)

Return type:

tuple[bool, int, ndarray[tuple[Any, ...], dtype[int_]] | None]

homodyne.optimization.nlsq.multistart.get_n_workers(config, n_starts)

Determine number of parallel workers.

Parameters:

• config (MultiStartConfig) – Multi-start configuration.

• n_starts (int) – Number of starting points.

Returns:

Number of workers to use.

Return type:

int

homodyne.optimization.nlsq.multistart.run_multistart_nlsq(data, bounds, config, single_fit_func, cost_func=None, custom_starts=None)

Run multi-start NLSQ optimization with FULL strategy.

NOTE: Only FULL strategy is supported. Subsampling is explicitly NOT used per project requirements - numerical precision takes priority over speed.

Parameters:

• data (dict[str, Any]) – XPCS data dictionary.

• bounds (ndarray[tuple[Any, ...], dtype[double]]) – Parameter bounds as (n_params, 2) array.

• config (MultiStartConfig) – Multi-start configuration.

• single_fit_func (Callable[[dict[str, Any], ndarray[tuple[Any, ...], dtype[double]]], SingleStartResult]) – Function that runs a single NLSQ fit. Signature: (data, initial_params) -> SingleStartResult

• cost_func (Callable[[ndarray[tuple[Any, ...], dtype[double]]], float] | None) – Function that computes cost for screening. Signature: (params) -> float

• custom_starts (list[list[float]] | ndarray[tuple[Any, ...], dtype[double]] | None) – User-provided custom starting points (overrides config.custom_starts).

Returns:

Aggregated results from all starting points.

Return type:

MultiStartResult

Key Classes

homodyne.optimization.nlsq.multistart.MultiStartConfig	Configuration for multi-start optimization.
homodyne.optimization.nlsq.multistart.MultiStartResult	Aggregated results from multi-start optimization.
homodyne.optimization.nlsq.multistart.SingleStartResult	Result from a single starting point optimization.

Configuration

Multi-start can be enabled via YAML configuration:

optimization:
  nlsq:
    multi_start:
      enable: true
      n_starts: 10
      seed: 42
      sampling_strategy: latin_hypercube
      use_screening: true
      screen_keep_fraction: 0.5
      refine_top_k: 3
      degeneracy_threshold: 0.1

Key Features

• Latin Hypercube Sampling: Better space-filling than random sampling

• Screening Phase: Filters poor starting points before expensive optimization

• Parallel Execution: Uses ProcessPoolExecutor for multi-core parallelism

• Basin Clustering: Identifies unique local minima in parameter space

• Degeneracy Detection: Warns when multiple solutions have similar chi-squared

• FULL Strategy Only: No subsampling per project requirements (numerical precision priority)

Determining n_starts

The number of starting points (n_starts) significantly impacts both optimization quality and computational cost. This section provides guidance for selecting appropriate values.

Minimum Requirements

For Latin Hypercube Sampling to provide adequate parameter space coverage, n_starts should be at least equal to the number of parameters:

Minimum n_starts by Analysis Mode

Analysis Mode	Parameters	Minimum n_starts
static_isotropic	5 (contrast, offset, D₀, α, D_offset)	5
laminar_flow	9 (+ γ̇₀, β, γ̇_offset, φ₀)	9
laminar_flow + per-angle (individual)	2×n_phi + 7	2×n_phi + 7
laminar_flow + per-angle (constant)	2 + 7 = 9	9

Impact of Anti-Degeneracy per_angle_mode

The per_angle_mode setting dramatically affects parameter count and thus n_starts:

Parameter Count by per_angle_mode (23-angle laminar_flow)

Mode	Per-Angle Params	Total Params	Recommended n_starts
individual	2 × 23 = 46	53	100-150
fourier (order=2)	2 × 5 = 10	17	20-40
constant	2	9	10-20

Constant mode (per_angle_mode: "constant") assumes all angles share the same contrast and offset, reducing parameter count from 53 to 9. This makes multi-start optimization tractable for many-angle datasets.

Recommended Settings by Use Case

n_starts Recommendations

Use Case	n_starts Formula	Description
Quick exploration	10	Default, fast baseline
Standard analysis	2 × n_params	Good coverage of parameter space
Degeneracy detection	3 × n_params	Better basin discovery
Publication quality	5 × n_params	Thorough exploration

Screening Considerations

When use_screening: true (default), only a fraction of starting points proceed to full optimization:

• With screen_keep_fraction: 0.5 (default): - 20 starts → 10 full optimizations - 100 starts → 50 full optimizations

Increase n_starts accordingly to achieve desired effective sample size.

Computational Cost

• Execution time scales linearly with effective n_starts

• For datasets ≥ 500K points: sequential execution (no parallelism benefit)

• Each fit runs the full optimization pipeline

Example Configuration

optimization:
  nlsq:
    # Use constant mode to reduce parameters (53 → 9)
    anti_degeneracy:
      enable: true
      per_angle_mode: "constant"
      constant_scaling_threshold: 3

    multi_start:
      enable: true
      n_starts: 20              # ~2× for 9 params
      use_screening: true
      screen_keep_fraction: 0.5 # 10 full fits
      seed: 42

Validation Warning

The code validates n_starts and warns if inadequate:

WARNING: n_starts (5) < n_params (9): LHS coverage may be inadequate.
Consider n_starts >= 9.

CLI Integration

homodyne --config config.yaml --method nlsq --multi-start

CMA-ES Global Optimization

CMA-ES (Covariance Matrix Adaptation Evolution Strategy) provides robust global optimization for multi-scale parameter estimation problems. It excels when parameter scales differ by several orders of magnitude, such as in laminar_flow mode (D₀ ~ 10⁴ vs γ̇₀ ~ 10⁻³).

CMA-ES global optimization wrapper for homodyne.

Provides CMA-ES integration using NLSQ’s CMAESOptimizer with: - Automatic memory configuration for large datasets - BIPOP restart strategy for robust convergence - Scale-ratio based method selection - Integration with homodyne’s model caching

CMA-ES (Covariance Matrix Adaptation Evolution Strategy) is particularly beneficial for XPCS laminar_flow mode where parameters have vastly different scales (e.g., D₀ ~ 1e4 vs γ̇₀ ~ 1e-3, scale ratio > 1e7).

NLSQ v0.6.4+ Features: - evosax backend for JAX-accelerated evolution - BIPOP restart strategy (alternating large/small populations) - Memory batching: population_batch_size, data_chunk_size - MethodSelector for auto-selection based on scale ratio

Usage

>>> from homodyne.optimization.nlsq.cmaes_wrapper import CMAESWrapper
>>> wrapper = CMAESWrapper()
>>> if wrapper.should_use_cmaes(bounds):
...     result = wrapper.fit(model_func, xdata, ydata, p0, bounds)

class homodyne.optimization.nlsq.cmaes_wrapper.CMAESWrapperConfig

Bases: object

Configuration for CMA-ES wrapper.

preset

CMA-ES preset: “cmaes-fast” (50 gen), “cmaes” (100 gen), “cmaes-global” (200 gen).

Type:

str

max_generations

Maximum CMA-ES generations. None = use preset default + adaptive scaling.

Type:

int | None

sigma

Initial step size as fraction of search range (0, 1].

Type:

float

tol_fun

Function value tolerance for convergence.

Type:

float

tol_x

Parameter tolerance for convergence.

Type:

float

restart_strategy

Restart strategy: “none” or “bipop”.

Type:

str

max_restarts

Maximum number of BIPOP restarts.

Type:

int

population_batch_size

Batch size for population evaluation (None = auto).

Type:

int | None

data_chunk_size

Chunk size for data streaming (None = auto).

Type:

int | None

refine_with_nlsq

Whether to refine CMA-ES solution with NLSQ TRF.

Type:

bool

auto_memory

Whether to auto-configure memory parameters.

Type:

bool

memory_limit_gb

Memory limit for auto-configuration in GB.

Type:

float

refinement_workflow

NLSQ workflow for refinement: “auto” (recommended), “standard”, “streaming”.

Type:

str

refinement_ftol

Function tolerance for NLSQ refinement.

Type:

float

refinement_xtol

Parameter tolerance for NLSQ refinement.

Type:

float

refinement_gtol

Gradient tolerance for NLSQ refinement.

Type:

float

refinement_max_nfev

Maximum function evaluations for NLSQ refinement.

Type:

int

refinement_loss

Loss function for NLSQ refinement: “linear”, “soft_l1”, “huber”, etc.

Type:

str

preset: str = 'cmaes'

max_generations: int | None = None

popsize: int | None = None

sigma: float = 0.5

sigma_warmstart: float = 0.05

tol_fun: float = 1e-08

tol_x: float = 1e-08

restart_strategy: str = 'bipop'

max_restarts: int = 9

population_batch_size: int | None = None

data_chunk_size: int | None = None

auto_memory: bool = True

memory_limit_gb: float = 8.0

refine_with_nlsq: bool = True

refinement_workflow: str = 'auto'

refinement_ftol: float = 1e-10

refinement_xtol: float = 1e-10

refinement_gtol: float = 1e-10

refinement_max_nfev: int = 500

refinement_loss: str = 'linear'

normalize: bool = True

normalization_epsilon: float = 1e-12

classmethod from_nlsq_config(config)

Create CMAESWrapperConfig from NLSQConfig.

Parameters:

config (NLSQConfig) – NLSQ configuration object.

Returns:

CMA-ES wrapper configuration.

Return type:

CMAESWrapperConfig

to_cmaes_config(n_params, *, sigma_override=None)

Convert to NLSQ CMAESConfig.

Parameters:

• n_params (int) – Number of parameters for popsize calculation.

• sigma_override (float | None) – If provided, override the default sigma value. Used to apply warm-start sigma when NLSQ warm-start is active.

Returns:

NLSQ CMAESConfig object.

Return type:

Any

Raises:

ImportError – If NLSQ CMA-ES is not available.

__init__(preset='cmaes', max_generations=None, popsize=None, sigma=0.5, sigma_warmstart=0.05, tol_fun=1e-08, tol_x=1e-08, restart_strategy='bipop', max_restarts=9, population_batch_size=None, data_chunk_size=None, auto_memory=True, memory_limit_gb=8.0, refine_with_nlsq=True, refinement_workflow='auto', refinement_ftol=1e-10, refinement_xtol=1e-10, refinement_gtol=1e-10, refinement_max_nfev=500, refinement_loss='linear', normalize=True, normalization_epsilon=1e-12)

class homodyne.optimization.nlsq.cmaes_wrapper.CMAESResult

Bases: object

Result from CMA-ES optimization.

parameters

Optimized parameter values.

Type:

np.ndarray

covariance

Parameter covariance matrix (if computed).

Type:

np.ndarray | None

chi_squared

Final chi-squared value.

Type:

float

success

Whether optimization converged successfully.

Type:

bool

diagnostics

CMA-ES diagnostics (generations, evaluations, etc.).

Type:

dict

method_used

Method used: “cmaes” or “multi-start”.

Type:

str

nlsq_refined

Whether result was refined with NLSQ L-M.

Type:

bool

message

Convergence message.

Type:

str

parameters: ndarray

covariance: ndarray | None

chi_squared: float

success: bool

diagnostics: dict

method_used: str = 'cmaes'

nlsq_refined: bool = False

message: str = ''

__init__(parameters, covariance, chi_squared, success, diagnostics=<factory>, method_used='cmaes', nlsq_refined=False, message='')

class homodyne.optimization.nlsq.cmaes_wrapper.CMAESWrapper

Bases: object

Wrapper around NLSQ’s CMAESOptimizer for homodyne integration.

This wrapper provides: - Scale-ratio based method selection (CMA-ES vs multi-start) - Automatic memory configuration for large datasets - BIPOP restart strategy for robust global optimization - Optional L-M refinement of CMA-ES solutions

Parameters:

config (CMAESWrapperConfig | None) – Configuration for CMA-ES wrapper. If None, uses defaults.

Examples

>>> wrapper = CMAESWrapper()
>>> if wrapper.should_use_cmaes(bounds, scale_threshold=1000):
...     result = wrapper.fit(model_func, xdata, ydata, p0, bounds)

__init__(config=None)

Initialize CMA-ES wrapper.

Parameters:

config (CMAESWrapperConfig | None) – Configuration for wrapper. Uses defaults if None.

property is_available: bool

Check if CMA-ES is available.

compute_scale_ratio(bounds)

Compute scale ratio from parameter bounds.

The scale ratio is the ratio of the largest to smallest parameter range. High scale ratios (> 1000) indicate multi-scale problems where CMA-ES excels.

Parameters:

bounds (tuple[ndarray, ndarray]) – Lower and upper bounds as (lower, upper) arrays.

Returns:

Scale ratio (max_range / min_range).

Return type:

float

Examples

>>> lower = np.array([0, 0.001, 100])
>>> upper = np.array([1, 0.01, 10000])
>>> wrapper.compute_scale_ratio((lower, upper))
11000.0  # (10000-100) / (0.01-0.001)

should_use_cmaes(bounds, scale_threshold=1000.0)

Determine if CMA-ES should be used based on scale ratio.

CMA-ES adapts its covariance matrix to different parameter scales, making it ideal for multi-scale optimization problems. This method checks if the scale ratio exceeds the threshold.

Parameters:

• bounds (tuple[ndarray, ndarray]) – Parameter bounds as (lower, upper) arrays.

• scale_threshold (float) – Scale ratio threshold for CMA-ES selection. Default: 1000.

Returns:

True if CMA-ES should be used.

Return type:

bool

Notes

XPCS laminar_flow mode typically has scale ratios > 1e7: - D₀ ~ 1e4 (diffusion coefficient) - γ̇₀ ~ 1e-3 (shear rate)

fit(model_func, xdata, ydata, p0, bounds, sigma=None, warmstart_chi2=None)

Run CMA-ES global optimization.

Parameters:

• model_func (Callable) – Model function: y = f(x, *params).

• xdata (ndarray) – Independent variable data.

• ydata (ndarray) – Dependent variable data to fit.

• p0 (ndarray) – Initial parameter guess.

• bounds (tuple[ndarray, ndarray]) – Parameter bounds as (lower, upper).

• sigma (ndarray | None) – Data uncertainties (optional).

• warmstart_chi2 (float | None) – Chi-squared from NLSQ warm-start. If provided and CMA-ES chi2 exceeds 10x this value, refinement is skipped (the comparison in core.py will discard the CMA-ES result anyway). Also triggers use of sigma_warmstart instead of sigma for the CMA-ES search.

Returns:

Optimization result with parameters, covariance, diagnostics.

Return type:

CMAESResult

Raises:

• ImportError – If CMA-ES is not available.

• RuntimeError – If optimization fails.

homodyne.optimization.nlsq.cmaes_wrapper.fit_with_cmaes(model_func, xdata, ydata, p0, bounds, sigma=None, config=None)

Convenience function for CMA-ES optimization.

Parameters:

• model_func (Callable) – Model function: y = f(x, *params).

• xdata (ndarray) – Independent variable data.

• ydata (ndarray) – Dependent variable data to fit.

• p0 (ndarray) – Initial parameter guess.

• bounds (tuple[ndarray, ndarray]) – Parameter bounds as (lower, upper).

• sigma (ndarray | None) – Data uncertainties (optional).

• config (CMAESWrapperConfig | None) – Configuration. Uses defaults if None.

Returns:

Optimization result.

Return type:

CMAESResult

Examples

>>> result = fit_with_cmaes(model, x, y, p0, bounds)
>>> print(f"Best params: {result.parameters}")

Key Classes

homodyne.optimization.nlsq.cmaes_wrapper.CMAESWrapper	Wrapper around NLSQ's CMAESOptimizer for homodyne integration.
homodyne.optimization.nlsq.cmaes_wrapper.CMAESWrapperConfig	Configuration for CMA-ES wrapper.
homodyne.optimization.nlsq.cmaes_wrapper.CMAESResult	Result from CMA-ES optimization.

When to Use CMA-ES

CMA-ES is recommended when:

• Multi-scale parameters: Scale ratio > 1000 (e.g., D₀/γ̇₀ > 10⁶)

• Complex loss landscapes: Multiple local minima, saddle points

• Poor initial guess: CMA-ES explores globally, not just locally

• laminar_flow mode: 7 physical parameters with vastly different scales

The CMAESWrapper.should_use_cmaes() method automatically detects multi-scale problems by computing the scale ratio from parameter bounds.

Two-Phase Architecture

Phase 1: CMA-ES Global Search
├─ Population-based evolutionary optimization
├─ Covariance matrix adapts to parameter scales
├─ BIPOP restart strategy (alternating large/small populations)
└─ Memory batching: population_batch_size, data_chunk_size

Phase 2: NLSQ TRF Refinement (if refine_with_nlsq=True)
├─ Uses NLSQ curve_fit with workflow="auto"
├─ Memory-aware: auto-selects standard/chunked/streaming
├─ Tighter tolerances (ftol=1e-10 vs CMA-ES 1e-8)
└─ Provides proper covariance matrix via Jacobian

Configuration

CMA-ES can be configured via YAML:

optimization:
  nlsq:
    cmaes:
      enable: true                      # Enable CMA-ES global optimization
      preset: "cmaes"                   # "cmaes-fast" (50), "cmaes" (100), "cmaes-global" (200)
      max_generations: 100              # Maximum CMA-ES generations
      sigma: 0.5                        # Initial step size (fraction of bounds)
      tol_fun: 1.0e-8                   # Function tolerance
      tol_x: 1.0e-8                     # Parameter tolerance
      restart_strategy: "bipop"         # "none" or "bipop"
      max_restarts: 9                   # Maximum BIPOP restarts
      population_batch_size: null       # null = auto, or explicit batch size
      data_chunk_size: null             # null = auto, or explicit chunk size
      refine_with_nlsq: true            # Refine with NLSQ TRF after CMA-ES
      auto_select: true                 # Auto-select CMA-ES vs multi-start
      scale_threshold: 1000.0           # Scale ratio threshold

      # NLSQ TRF Refinement Settings
      refinement_workflow: "auto"       # "auto", "standard", "streaming"
      refinement_ftol: 1.0e-10          # Tighter for local refinement
      refinement_xtol: 1.0e-10
      refinement_gtol: 1.0e-10
      refinement_max_nfev: 500          # Bounded iterations
      refinement_loss: "linear"         # "linear", "soft_l1", "huber"

Usage Example

from homodyne.optimization.nlsq.cmaes_wrapper import CMAESWrapper, CMAESWrapperConfig

# Create wrapper with custom config
wrapper = CMAESWrapper(CMAESWrapperConfig(
    preset="cmaes",
    refine_with_nlsq=True,
    refinement_ftol=1e-10,
))

# Check if CMA-ES is appropriate for this problem
if wrapper.should_use_cmaes(bounds, scale_threshold=1000):
    result = wrapper.fit(model_func, xdata, ydata, p0, bounds)
    print(f"Chi²: {result.chi_squared:.4e}")
    print(f"Refined: {result.nlsq_refined}")
    print(f"Improvement: {result.diagnostics.get('chi_squared_improvement', 0):.2%}")

CMA-ES vs Multi-Start vs Local

Method	Best For	Convergence	Memory
CMA-ES	Multi-scale (ratio > 1000)	Global (covariance)	Bounded
Multi-start	Multiple local minima	Local from N starts	N × single fit
Local (TRF)	Good initial guess	Local (quadratic)	Jacobian-based

Streaming Optimizer for Large Datasets

For datasets exceeding available memory (>10M points on typical systems), the NLSQ wrapper automatically switches to streaming optimization using mini-batch gradient descent. This eliminates OOM errors by processing data in small batches.

Why Streaming?

Standard Levenberg-Marquardt optimization requires computing a dense Jacobian matrix (n_points × n_params × 8 bytes) plus JAX autodiff intermediates (~3× Jacobian size). For 23M points with 53 parameters, this exceeds 30 GB. Streaming mode processes data in 10K-point batches, keeping memory usage below 2 GB.

Memory-Based Auto-Selection

The NLSQWrapper._should_use_streaming() method estimates peak memory usage and automatically selects streaming mode when:

1. Estimated memory > memory_threshold_gb (default: 16 GB), OR

2. Estimated memory > 70% of available system RAM

Decision Logic:

fit() called
      │
      ▼
┌─────────────────────────────────────────┐
│ Estimate memory for Jacobian + autodiff │
│ = n_points × n_params × 8 × 4           │
└─────────────────────────────────────────┘
      │
      ▼
┌─────────────────────────────────────────┐
│ Estimated > threshold OR > 70% RAM?     │
└─────────────────────────────────────────┘
      │                    │
     YES                   NO
      │                    │
      ▼                    ▼
┌─────────────┐     ┌─────────────────────┐
│ Streaming   │     │ Stratified L-M      │
│ Optimizer   │     │ (Full Jacobian)     │
│             │     │                     │
│ Mini-batch  │     │ Trust-region        │
│ L-BFGS      │     │ Levenberg-Marquardt │
└─────────────┘     └─────────────────────┘

Configuration

Streaming mode can be configured via YAML:

optimization:
  nlsq:
    # Memory threshold for automatic streaming mode (GB)
    memory_threshold_gb: 16.0

    # Force streaming mode regardless of memory (default: false)
    use_streaming: false

    # Streaming optimizer settings
    streaming:
      batch_size: 10000       # Points per mini-batch
      max_epochs: 50          # Maximum training epochs
      learning_rate: 0.001    # L-BFGS line search scale
      convergence_tol: 1e-6   # Convergence tolerance

Performance Characteristics

Mode	Memory Usage	Convergence	Time (23M points)
Stratified L-M	~30+ GB	Exact (Newton)	10-15 min
Streaming	~2 GB	Approximate (L-BFGS)	15-30 min

When to Use:

• Stratified L-M (default): Datasets < 10M points, sufficient RAM (64GB+)

• Streaming: Datasets > 10M points, memory-constrained systems (32GB RAM)

Implementation Details

The streaming optimizer uses NLSQ’s AdaptiveHybridStreamingOptimizer class:

from nlsq import AdaptiveHybridStreamingOptimizer, HybridStreamingConfig

config = HybridStreamingConfig(
    chunk_size=50000,
    warmup_iterations=100,
    max_warmup_iterations=500,
    gauss_newton_max_iterations=50,
    gauss_newton_tol=1e-8,
    normalize=True,
    normalization_strategy="bounds",
)

optimizer = AdaptiveHybridStreamingOptimizer(config)
result = optimizer.fit(
    data_source=(x_data, y_data),
    func=model_fn,
    p0=initial_params,
    bounds=bounds,
)

Key features:

• 4-phase hybrid optimization: L-BFGS warmup + Gauss-Newton refinement

• Parameter normalization: Equalizes gradient magnitudes across multi-scale parameters

• Exact J^T J accumulation: Proper covariance estimation in streaming mode

• Chunk-based processing: Memory-efficient for unlimited dataset sizes

• Progress tracking: Logs phase progress and convergence metrics

CMC: Consensus Monte Carlo

CMC provides Bayesian parameter estimation with full posterior sampling using NumPyro/NUTS.

Key Features:

• Physics-informed priors

• Automatic retry mechanism (max 3 attempts)

• Single-angle log-space D0 sampling for stability

• ArviZ-native output format

Core Module

CMC core module - main entry point.

This module provides the fit_mcmc_jax() function that serves as the main entry point for CMC analysis, matching the CLI signature.

homodyne.optimization.cmc.core.fit_mcmc_jax(data, t1, t2, phi, q, L, analysis_mode, method='mcmc', cmc_config=None, initial_values=None, parameter_space=None, dt=None, output_dir=None, progress_bar=True, run_id=None, nlsq_result=None, **kwargs)

Run CMC (Consensus Monte Carlo) analysis on XPCS data.

This function signature matches the CLI call in cli/commands.py:1201.

Parameters:

• data (ndarray) – Pooled C2 correlation data, shape (n_total,).

• t1 (ndarray) – Pooled time coordinates t1, shape (n_total,).

• t2 (ndarray) – Pooled time coordinates t2, shape (n_total,).

• phi (ndarray) – Pooled phi angles, shape (n_total,).

• q (float) – Wavevector magnitude.

• L (float) – Stator-rotor gap length (nm).

• analysis_mode (str) – Analysis mode: “static” or “laminar_flow”.

• method (str) – Method identifier (always “mcmc” for CMC).

• cmc_config (dict[str, Any] | None) – CMC configuration from ConfigManager.get_cmc_config().

• initial_values (dict[str, float] | None) – Initial parameter values from ConfigManager.get_initial_parameters().

• parameter_space (ParameterSpace | None) – Parameter space with bounds and priors from ParameterSpace.from_config().

• dt (float | None) – Time step for physics model. If None, inferred from pooled time arrays.

• output_dir (Path | str | None) – Output directory for saving results.

• progress_bar (bool) – Whether to show progress bar during sampling.

• run_id (str | None) – Optional identifier used to correlate logs across shards/backends.

• nlsq_result (dict | None) – Optional NLSQ result dictionary for warm-start priors. When provided, builds informative priors centered on NLSQ estimates, improving convergence speed and reducing divergences. Should contain parameter values and optionally uncertainties (see extract_nlsq_values_for_cmc).

• **kwargs – Additional keyword arguments (for compatibility).

Returns:

Complete result with posterior samples and diagnostics.

Return type:

CMCResult

Raises:

• ValueError – If data validation fails.

• RuntimeError – If MCMC sampling fails.

Examples

>>> from homodyne.optimization.cmc import fit_mcmc_jax
>>> result = fit_mcmc_jax(
...     data=c2_pooled,
...     t1=t1_pooled,
...     t2=t2_pooled,
...     phi=phi_pooled,
...     q=0.01,
...     L=2000000.0,
...     analysis_mode="laminar_flow",
...     method="mcmc",
...     cmc_config=config.get_cmc_config(),
...     initial_values=config.get_initial_parameters(),
...     parameter_space=parameter_space,
... )
>>> print(result.convergence_status)
converged

homodyne.optimization.cmc.core.run_cmc_analysis(data, t1, t2, phi, q, L, analysis_mode, config, parameter_space, initial_values=None, dt=None)

Simplified interface for CMC analysis.

This is a convenience wrapper around fit_mcmc_jax() that takes a CMCConfig object directly instead of a dict.

Parameters:

• data (ndarray) – Data arrays.

• t1 (ndarray) – Data arrays.

• t2 (ndarray) – Data arrays.

• phi (ndarray) – Data arrays.

• q (float) – Physics parameters.

• L (float) – Physics parameters.

• analysis_mode (str) – Analysis mode.

• config (CMCConfig) – CMC configuration object.

• parameter_space (ParameterSpace) – Parameter space.

• initial_values (dict[str, float] | None) – Initial values.

• dt (float | None) – Time step (None infers from pooled time arrays).

Returns:

Analysis result.

Return type:

CMCResult

Configuration

CMC configuration dataclass and validation.

This module provides the CMCConfig dataclass for parsing and validating CMC-specific configuration settings from the YAML config file.

Config Precedence (Important)

The CLI reads base optimization.mcmc settings and applies them to per_shard_mcmc. This means if base mcmc differs from per_shard_mcmc in your YAML config, the CLI will overwrite per_shard_mcmc with base values. To avoid surprises, keep base mcmc and per_shard_mcmc aligned.

Example aligned config:

optimization:
  mcmc:
    num_warmup: 500
    num_samples: 1500
    num_chains: 4
  cmc:
    per_shard_mcmc:
      num_warmup: 500
      num_samples: 1500
      num_chains: 4

class homodyne.optimization.cmc.config.CMCConfig

Bases: object

Configuration for Consensus Monte Carlo (CMC) analysis.

enable

Whether to enable CMC. “auto” enables based on data size.

Type:

bool | str

min_points_for_cmc

Minimum data points to trigger CMC mode.

Type:

int

sharding_strategy

How to partition data: “stratified”, “random”, “contiguous”.

Type:

str

num_shards

Number of data shards. “auto” calculates from data size.

Type:

int | str

max_points_per_shard

Maximum points per shard. “auto” calculates optimally based on dataset size, analysis mode, and angle count (see _resolve_max_points_per_shard). Default: “auto”. Typical auto values: 5–20K for laminar_flow, 10–20K for static (scales with dataset size).

Type:

int | str

backend_name

Execution backend: “auto”, “multiprocessing”, “pjit”, “pbs”.

Type:

str

enable_checkpoints

Whether to save checkpoints during sampling.

Type:

bool

checkpoint_dir

Directory for checkpoint files.

Type:

str

num_warmup

Number of warmup/burn-in samples per chain.

Type:

int

num_samples

Number of posterior samples per chain.

Type:

int

num_chains

Number of MCMC chains.

Type:

int

chain_method

MCMC chain execution method. "parallel" (default) runs chains concurrently via JAX vectorization. "sequential" runs chains one at a time. Parallel is faster on multi-core CPUs but adds ~5-15% overhead on very small shards (<500 points); the sampler auto-falls-back to sequential in that case.

Type:

str

target_accept_prob

Target acceptance probability for NUTS.

Type:

float

dense_mass

Use dense mass matrix for NUTS. When True, learns parameter correlations for more efficient sampling. Default: True.

Type:

bool

max_r_hat

Maximum R-hat for convergence.

Type:

float

min_ess

Minimum effective sample size.

Type:

float

combination_method

How to combine shard posteriors. Options:

• "consensus_mc": Correct Consensus Monte Carlo (precision-weighted means). Recommended. Combines per-shard posterior moments, then generates new samples from the combined Gaussian.

• "weighted_gaussian": Legacy element-wise weighted averaging (deprecated).

• "simple_average": Simple element-wise averaging (deprecated).

Type:

str

min_success_rate

Minimum fraction of shards that must succeed.

Type:

float

run_id

Optional identifier used for structured logging across shards.

Type:

str | None

per_angle_mode

Per-angle scaling mode for anti-degeneracy defense (v2.18.0+):

• "auto": Auto-selects based on n_phi threshold (recommended). When n_phi >= threshold: Estimates per-angle values, AVERAGES them, broadcasts single value to all angles (matches NLSQ behavior). When n_phi < threshold: Uses individual mode.

• "constant": Per-angle contrast/offset from quantile estimation, used DIRECTLY (different fixed value per angle, NOT averaged). Reduces to 8 params (7 physical + 1 sigma).

• "individual": Independent contrast + offset per angle, all sampled. May suffer from parameter absorption degeneracy with many angles.

Type:

str

constant_scaling_threshold

n_phi threshold for auto mode’s per-angle strategy. When n_phi >= threshold, auto mode samples averaged contrast/offset (single value broadcast to all angles). When n_phi < threshold, auto mode falls back to individual per-angle sampling. Default: 3.

Type:

int

enable: bool | str = 'auto'

min_points_for_cmc: int = 100000

per_angle_mode: str = 'auto'

constant_scaling_threshold: int = 3

sharding_strategy: str = 'random'

num_shards: int | str = 'auto'

max_points_per_shard: int | str = 'auto'

backend_name: str = 'auto'

enable_checkpoints: bool = True

checkpoint_dir: str = './checkpoints/cmc'

num_warmup: int = 500

num_samples: int = 1500

num_chains: int = 4

chain_method: str = 'parallel'

target_accept_prob: float = 0.85

dense_mass: bool = True

adaptive_sampling: bool = True

max_tree_depth: int = 10

min_warmup: int = 100

min_samples: int = 200

enable_jax_profiling: bool = False

jax_profile_dir: str = './profiles/jax'

max_r_hat: float = 1.1

min_ess: float = 400.0

max_divergence_rate: float = 0.1

combination_method: str = 'robust_consensus_mc'

min_success_rate: float = 0.9

run_id: str | None = None

per_shard_timeout: int = 3600

heartbeat_timeout: int = 600

min_success_rate_warning: float = 0.8

require_nlsq_warmstart: bool = False

use_nlsq_informed_priors: bool = True

nlsq_prior_width_factor: float = 2.0

prior_tempering: bool = True

max_parameter_cv: float = 1.0

heterogeneity_abort: bool = True

min_points_per_shard: int = 10000

min_points_per_param: int = 1500

reparameterization_d_total: bool = True

reparameterization_log_gamma: bool = True

bimodal_min_weight: float = 0.2

bimodal_min_separation: float = 0.5

seed: int = 42

classmethod from_dict(config_dict)

Create CMCConfig from configuration dictionary.

Parameters:

config_dict (dict[str, Any]) – CMC configuration dictionary from ConfigManager.get_cmc_config().

Returns:

Validated configuration object.

Return type:

CMCConfig

Raises:

ValueError – If required fields are missing or invalid.

validate()

Validate configuration values.

Returns:

List of validation error messages (empty if valid).

Return type:

list[str]

is_valid()

Check if configuration is valid.

Returns:

True if configuration has no validation errors.

Return type:

bool

should_enable_cmc(n_points, analysis_mode=None)

Determine if CMC should be enabled for given data size.

Parameters:

• n_points (int) – Number of data points.

• analysis_mode (str | None) – Deprecated — ignored. Kept for backward compatibility.

Returns:

True if CMC should be enabled.

Return type:

bool

Notes

Threshold is min_points_for_cmc (default 100,000) for all modes.

get_num_shards(n_points, n_phi, n_params=7)

Calculate number of shards with param-aware sizing.

Parameters:

• n_points (int) – Total number of data points.

• n_phi (int) – Number of phi angles.

• n_params (int) – Number of model parameters (default: 7 for static).

Returns:

Number of shards to use.

Return type:

int

get_adaptive_sample_counts(shard_size, n_params=7)

Calculate adaptive warmup/samples based on shard size.

Small datasets benefit from fewer NUTS samples because: 1. JIT compilation overhead is amortized over fewer samples 2. Step size adaptation converges faster with simple likelihoods 3. Mass matrix estimation requires fewer warmup iterations

Profiling showed 1310s for 50 points with 500 warmup + 1500 samples. Adaptive scaling reduces this by 60-80% while maintaining statistical validity (ESS targets are reduced proportionally).

Parameters:

• shard_size (int) – Number of data points in the shard.

• n_params (int) – Number of model parameters (affects minimum samples).

Returns:

(num_warmup, num_samples) adjusted for shard size.

Return type:

tuple[int, int]

get_effective_per_angle_mode(n_phi, nlsq_per_angle_mode=None, has_nlsq_warmstart=False)

Determine effective per-angle mode based on configuration and data.

Parameters:

• n_phi (int) – Number of phi angles in the dataset.

• nlsq_per_angle_mode (str | None) – Optional per-angle mode from NLSQ result. When provided (from warm-start), CMC will use this mode to ensure parameterization parity with NLSQ. This prevents CMC vs NLSQ divergence from different model structures.

• has_nlsq_warmstart (bool) – Whether an NLSQ warm-start result is available. When True and both CMC and NLSQ use “auto” mode, upgrades to “constant_averaged” for fewer sampled parameters and better stability.

Returns:

Effective mode: “auto”, “constant”, “constant_averaged”, or “individual”.

Return type:

str

Notes

Mode semantics (same as NLSQ):

• auto: Sample single averaged contrast/offset (10 params for laminar_flow). Only activated when n_phi >= threshold (many angles).

• constant: Use FIXED per-angle values from quantile estimation (8 params).

• constant_averaged: Use FIXED averaged scaling for NLSQ parity.

• individual: Sample per-angle contrast/offset (n_phi*2 + 7 + 1 params).

Priority: nlsq_per_angle_mode > explicit config > auto-selection

When NLSQ warm-start is present and both sides use “auto”, upgrades to “constant_averaged” to fix scaling values and reduce parameter count. This prevents contrast/offset sampling from absorbing physical parameter signal, which was the root cause of heterogeneous shard posteriors.

to_dict()

Convert configuration to dictionary.

Returns:

Configuration as dictionary.

Return type:

dict[str, Any]

__init__(enable='auto', min_points_for_cmc=100000, per_angle_mode='auto', constant_scaling_threshold=3, sharding_strategy='random', num_shards='auto', max_points_per_shard='auto', backend_name='auto', enable_checkpoints=True, checkpoint_dir='./checkpoints/cmc', num_warmup=500, num_samples=1500, num_chains=4, chain_method='parallel', target_accept_prob=0.85, dense_mass=True, adaptive_sampling=True, max_tree_depth=10, min_warmup=100, min_samples=200, enable_jax_profiling=False, jax_profile_dir='./profiles/jax', max_r_hat=1.1, min_ess=400.0, max_divergence_rate=0.1, combination_method='robust_consensus_mc', min_success_rate=0.9, run_id=None, per_shard_timeout=3600, heartbeat_timeout=600, min_success_rate_warning=0.8, require_nlsq_warmstart=False, use_nlsq_informed_priors=True, nlsq_prior_width_factor=2.0, prior_tempering=True, max_parameter_cv=1.0, heterogeneity_abort=True, min_points_per_shard=10000, min_points_per_param=1500, reparameterization_d_total=True, reparameterization_log_gamma=True, bimodal_min_weight=0.2, bimodal_min_separation=0.5, seed=42, _validation_errors=<factory>)

CMC vs Single-Shard MCMC Decision Logic

CMC uses a unified sampler architecture: both single-shard (standard) MCMC and per-shard CMC sampling use the identical run_nuts_sampling() function from homodyne/optimization/cmc/sampler.py. The only difference is data volume and orchestration:

Decision Flow (see homodyne/optimization/cmc/core.py:620-664):

fit_mcmc_jax() called
      │
      ▼
┌─────────────────────────────────────────┐
│ n_points >= min_points_for_cmc (500K)?  │
│         OR explicit shards requested?   │
└─────────────────────────────────────────┘
      │                    │
     YES                   NO
      │                    │
      ▼                    ▼
┌─────────────┐     ┌─────────────────────┐
│ CMC Path    │     │ Single-Shard Path   │
│             │     │                     │
│ 1. Shard    │     │ run_nuts_sampling() │
│    data     │     │ with ALL data       │
│             │     │                     │
│ 2. Backend  │     │ Returns MCMCSamples │
│    runs     │     │ directly            │
│    run_nuts │     └─────────────────────┘
│    _sampling│
│    per shard│
│             │
│ 3. Combine  │
│    posteriors│
└─────────────┘

Comparison:

Aspect	Single-Shard (Standard)	CMC (Sharded)
Data handling	All points in one call	Subsets per shard (e.g., 10K each)
Execution	Single run_nuts_sampling()	Backend orchestrates parallel run_nuts_sampling() per shard
Results	Direct posterior samples	Combined via precision-weighted Gaussian consensus
Parallelization	Within-chain only	Across shards + within-chain
Memory	Must fit entire dataset	Each shard fits independently
Typical use	< 500K points	> 500K points

Key Configuration Parameter:

The min_points_for_cmc threshold (default: 500,000) controls automatic switching:

optimization:
  cmc:
    enable: "auto"              # "auto" | true | false
    min_points_for_cmc: 500000  # Threshold for auto-enable

• enable: "auto": Uses CMC when n_points >= min_points_for_cmc

• enable: true: Always uses CMC sharding (even for small datasets)

• enable: false: Always uses single-shard MCMC

Code Reference:

The decision is made in fit_mcmc_jax() (core.py:425-508):

# Determine if CMC sharding is needed
use_cmc = config.should_enable_cmc(prepared.n_total) or forced_shards

if shards is not None and len(shards) > 1:
    # CMC path: parallel backend
    backend = select_backend(config)
    mcmc_samples = backend.run(
        model=xpcs_model_scaled,
        ...
    )
else:
    # Single-shard path: direct sampling
    mcmc_samples, stats = run_nuts_sampling(
        model=xpcs_model_scaled,
        model_kwargs=model_kwargs,
        config=config,
        ...
    )

Both paths use identical:

• Model: xpcs_model_scaled (scaled/z-space parameterization)

• Sampler: run_nuts_sampling() with NumPyro NUTS

• Configuration: num_warmup, num_samples, num_chains, target_accept_prob

• Gradient balancing: Dense mass matrix (dense_mass=True)

Sharding Strategy (Detailed)

CMC (Consensus Monte Carlo) partitions large datasets into smaller shards that can be processed in parallel. Each shard runs independent NUTS sampling, and posteriors are combined using weighted Gaussian consensus.

Why Sharding?

NUTS MCMC is O(n) per iteration—it evaluates ALL data points in a shard for each gradient computation. For XPCS datasets with millions of points, a single NUTS run would take days. Sharding enables:

1. Parallelization: Run multiple shards simultaneously across CPU cores

2. Memory efficiency: Each shard fits in available RAM

3. Timeout management: Per-shard timeouts prevent runaway computations

Shard Size Selection Algorithm

The _resolve_max_points_per_shard() function automatically selects optimal shard sizes based on analysis mode and dataset size:

Laminar Flow Mode (7 parameters, complex gradients):

Dataset Size	Shard Size	Est. Shards	Per-Shard Runtime
< 2M points	4,800	~400	~1-2 min
2M - 50M points	3,000	600-16K	~1 min
50M - 100M points	4,800	10K-20K	~1 min
100M - 1B points	4,800	20K-50K	<1 min
1B+ points	6,000-10,000	100K+	<1 min

Static Mode (3 parameters, simpler gradients):

Dataset Size	Shard Size	Est. Shards	Per-Shard Runtime
< 50M points	100,000	~500	~5-10 min
50M - 100M points	80,000	~1K	~5 min
100M+ points	50,000	~2K+	~3-5 min

Key insight: Laminar flow uses ~20x smaller shards than static mode. The reparameterization to (D_ref, gamma_ref) produces unimodal posteriors, enabling 3-5K shards with adaptive sampling and prior tempering.

Sharding Strategies

CMC supports two sharding strategies:

Stratified Sharding (default, recommended)

Partitions data by phi angle. Each shard contains data for one angle:

• Preserves physical grouping of measurements

• Enables per-angle posterior estimates

• If an angle exceeds max_points_per_shard, it’s split into multiple shards

• Cap: max_shards_per_angle=100 (increases shard size if exceeded)

shards = shard_data_stratified(
    prepared,
    num_shards=None,  # Auto-calculate
    max_points_per_shard=5000,  # For laminar_flow
    max_shards_per_angle=100,
)

Random Sharding

Used when there’s only one phi angle but the dataset is large:

• Shuffles data indices randomly

• Splits into approximately equal parts

• Sorts within each shard to preserve temporal structure

• ALL data is used (no subsampling)

shards = shard_data_random(
    prepared,
    num_shards=None,
    max_points_per_shard=10000,
    max_shards=100,  # Cap to prevent memory issues
)

Memory Scalability

Each shard result contains posterior samples that must be held in memory during combination. Memory requirements scale with shard count:

Platform	Available RAM	Max Shards	Max Dataset (laminar)
Personal workstation	~20 GB	~500	~5M points
Bebop (36 cores)	~100 GB	~2,500	~25M points
Improv (128 cores)	~200 GB	~5,000	~50M points

Memory formula: Each shard result ≈ 100 KB (13 params × 2 chains × 1500 samples × 8 bytes). Peak memory ≈ 6 × K MB where K = number of shards.

The algorithm automatically caps shard count (default: 2000) and increases shard size to prevent memory exhaustion. For very large datasets exceeding limits, a warning is logged.

Runtime Estimation

CMC provides runtime estimates before sampling begins:

Runtime estimate: 2.5h total (100 shards / 18 workers, ~15min/shard with 4000 iterations)

The estimate accounts for:

• JIT compilation overhead (~30-60s per worker)

• MCMC iterations: num_chains × (num_warmup + num_samples)

• Points per shard and analysis mode complexity

• Parallel execution across available workers

After completion, actual vs. estimated runtime is logged with recommendations:

Runtime: 2.1h actual vs 2.5h estimated (84% - close to estimate)

Configuration Reference

Full YAML configuration for sharding:

optimization:
  cmc:
    enable: auto  # true, false, or "auto" (based on data size)
    min_points_for_cmc: 500000  # Threshold for auto-enable

    sharding:
      strategy: stratified  # "stratified" or "random"
      num_shards: auto  # "auto" or explicit integer
      max_points_per_shard: auto  # "auto" or explicit integer

    backend_config:
      name: multiprocessing  # "auto", "multiprocessing", "pjit", "pbs"
      enable_checkpoints: true
      checkpoint_dir: ./checkpoints/cmc

    per_shard_mcmc:
      num_warmup: 500
      num_samples: 1500
      num_chains: 2
      target_accept_prob: 0.85
      # Adaptive Sampling
      adaptive_sampling: true           # Scale by shard size
      max_tree_depth: 10                # NUTS tree depth limit
      min_warmup: 100                   # Minimum warmup floor
      min_samples: 200                  # Minimum samples floor
      # JAX Profiling
      enable_jax_profiling: false       # XLA-level profiling
      jax_profile_dir: "./profiles/jax"

    validation:
      max_per_shard_rhat: 1.1
      min_per_shard_ess: 100

    combination:
      method: robust_consensus_mc  # MAD-based outlier detection (default)
      min_success_rate: 0.90

    validation:
      max_per_shard_rhat: 1.1
      min_per_shard_ess: 100
      max_divergence_rate: 0.10       # Quality filter: exclude shards >10%
      require_nlsq_warmstart: false   # Require NLSQ warm-start

    per_shard_timeout: 3600  # 1 hour max per shard (reduced)
    heartbeat_timeout: 600   # 10 min worker heartbeat

Critical settings for laminar_flow:

• Use max_points_per_shard: auto (resolves to 3K-5K based on size)

• Do NOT set max_points_per_shard: 100000 — this causes 1-2+ hour per-shard runtimes

• Keep num_warmup and num_samples aligned between mcmc and per_shard_mcmc

• Consider require_nlsq_warmstart: true for production runs (reduces divergences from ~28% to <5%)

Quality Filtering:

The max_divergence_rate setting automatically filters out shards with excessive divergent transitions before consensus combination:

optimization:
  cmc:
    validation:
      max_divergence_rate: 0.10  # Exclude shards with >10% divergence

Shards with divergence rate exceeding this threshold are excluded from the final posterior combination, preventing corrupted posteriors from biasing estimates.

NLSQ Warm-Start Requirement:

For laminar_flow mode with 7 parameters spanning 6+ orders of magnitude, cold-start CMC runs often show high divergence rates (28%+) and inflated uncertainty. Enable warm-start requirement for production:

optimization:
  cmc:
    validation:
      require_nlsq_warmstart: true

When enabled, fit_mcmc_jax() will raise ValueError if called without nlsq_result or initial_values for laminar_flow mode

Adaptive Sampling:

Adaptive sampling automatically scales warmup and sample counts based on shard size, reducing NUTS overhead by 60-80% for small datasets while maintaining statistical validity.

optimization:
  cmc:
    per_shard_mcmc:
      adaptive_sampling: true     # Enable adaptive scaling
      max_tree_depth: 10          # Limit NUTS tree depth (2^10 max leapfrog)
      min_warmup: 100             # Floor for warmup scaling
      min_samples: 200            # Floor for samples scaling

The scaling formula uses a 10K point reference:

• scale_factor = min(1.0, shard_size / 10000)

• Small shards (< 10K points) get proportionally fewer warmup/samples

• Minimum samples scale with parameter count: max(min_samples, 50 × n_params)

This optimization was informed by profiling showing that XLA JIT compilation and NUTS leapfrog integration dominate runtime (not Python overhead), making sample count reduction the most effective optimization.

JAX Profiling:

XLA-level profiling for diagnosing NUTS performance bottlenecks. Standard Python profilers (py-spy, cProfile) cannot see inside JIT-compiled code.

optimization:
  cmc:
    per_shard_mcmc:
      enable_jax_profiling: true
      jax_profile_dir: "./profiles/jax"

View profiles with TensorBoard: tensorboard --logdir=./profiles/jax

Practical Guidelines

For typical 3-angle, 3M point laminar_flow datasets:

optimization:
  cmc:
    sharding:
      max_points_per_shard: auto  # Will select ~10K-20K
    per_shard_mcmc:
      num_warmup: 300
      num_samples: 700
      num_chains: 2

Expected: ~150-300 shards, ~5-8 min/shard, ~2-4 hours total on 18-core workstation.

For 50M+ point production datasets on HPC:

optimization:
  cmc:
    sharding:
      max_points_per_shard: auto  # Will select ~6K-8K
    per_shard_mcmc:
      num_warmup: 500
      num_samples: 1500
      num_chains: 2
    per_shard_timeout: 7200  # 2 hours

Expected: ~6K-8K shards, parallel execution across cluster nodes.

Model Definition

NumPyro model definition for MCMC sampling.

NumPyro model for XPCS C2 correlation function.

This module defines the probabilistic model for Bayesian inference of XPCS parameters using NumPyro.

CRITICAL: Parameter sampling order must match: 1. Per-angle contrast: contrast_0, contrast_1, … (individual mode only) 2. Per-angle offset: offset_0, offset_1, … (individual mode only) 3. Physical parameters: D0, alpha, D_offset, [gamma_dot_t0, …]

Per-Angle Modes (v2.18.0+): - “individual”: Independent contrast + offset per angle (2*n_phi + n_physical + 1 params) - “constant”: Fixed per-angle contrast/offset from quantile estimation (n_physical + 1 params) - “auto”: Selects based on n_phi threshold (constant if n_phi >= 3, else individual)

homodyne.optimization.cmc.model.validate_model_output(c2_theory, params)

Validate that model output is physically reasonable.

Parameters:

• c2_theory (Array) – Theoretical C2 values.

• params (Array) – Parameter values.

Returns:

True if output is valid.

Return type:

bool | Array

homodyne.optimization.cmc.model.get_model_param_count(n_phi, analysis_mode, per_angle_mode='individual')

Get total number of sampled parameters.

Parameters:

• n_phi (int) – Number of phi angles.

• analysis_mode (str) – Analysis mode.

• per_angle_mode (str) – Per-angle scaling mode: “individual”, “auto”, or “constant”.

Returns:

Total number of parameters (including sigma).

Return type:

int

Notes

Mode semantics (same as NLSQ): - individual mode: 2*n_phi (contrast + offset) + physical + sigma - auto mode: 2 (averaged contrast + offset, SAMPLED) + physical + sigma - constant mode: 0 per-angle (FIXED from quantiles) + physical + sigma

homodyne.optimization.cmc.model.xpcs_model_scaled(data, t1, t2, phi_unique, phi_indices, q, L, dt, analysis_mode, parameter_space, n_phi, time_grid=None, noise_scale=0.1, num_shards=1, shard_grid=None, **kwargs)

NumPyro model with non-centered parameterization for gradient balancing.

This model samples all parameters in normalized (z) space where z ~ N(0,1), then transforms to original space: P = center + scale * z. This ensures all gradient magnitudes are balanced, solving the 0% acceptance rate issue caused by D0 (~10^4) dominating gradients over gamma_dot_t0 (~10^-3).

The physics computation is identical to xpcs_model, only the sampling space is transformed.

Parameters:

• data (Array) – Observed C2 correlation data, shape (n_total,).

• t1 (Array) – Time coordinates, shape (n_total,).

• t2 (Array) – Time coordinates, shape (n_total,).

• phi_unique (Array) – Unique phi angles, shape (n_phi,).

• phi_indices (Array) – Index into per-angle arrays for each point, shape (n_total,).

• q (float) – Wavevector magnitude.

• L (float) – Stator-rotor gap length (nm).

• dt (float) – Time step.

• analysis_mode (str) – Analysis mode: “static” or “laminar_flow”.

• parameter_space (ParameterSpace) – Parameter space with bounds and priors.

• n_phi (int) – Number of unique phi angles.

• noise_scale (float) – Initial estimate of observation noise.

Return type:

None

homodyne.optimization.cmc.model.xpcs_model_constant(data, t1, t2, phi_unique, phi_indices, q, L, dt, analysis_mode, parameter_space, n_phi, time_grid=None, noise_scale=0.1, fixed_contrast=None, fixed_offset=None, num_shards=1, shard_grid=None, **kwargs)

NumPyro model with FIXED per-angle scaling (anti-degeneracy constant mode).

This model uses FIXED per-angle contrast/offset values estimated from quantile analysis of the raw data. These values are NOT sampled, reducing the parameter space to only physical parameters + sigma.

This matches NLSQ’s anti-degeneracy constant mode and prevents parameter absorption degeneracy where per-angle params absorb physical signals.

Parameter count comparison (laminar_flow, n_phi=23): - individual mode: 54 params (46 per-angle + 7 physical + 1 sigma) - constant mode: 8 params (7 physical + 1 sigma)

Parameters:

• data (Array) – Observed C2 correlation data, shape (n_total,).

• t1 (Array) – Time coordinates, shape (n_total,).

• t2 (Array) – Time coordinates, shape (n_total,).

• phi_unique (Array) – Unique phi angles, shape (n_phi,).

• phi_indices (Array) – Index into per-angle arrays for each point, shape (n_total,).

• q (float) – Wavevector magnitude.

• L (float) – Stator-rotor gap length (nm).

• dt (float) – Time step.

• analysis_mode (str) – Analysis mode: “static” or “laminar_flow”.

• parameter_space (ParameterSpace) – Parameter space with bounds and priors.

• n_phi (int) – Number of unique phi angles.

• noise_scale (float) – Initial estimate of observation noise.

• fixed_contrast (Array | None) – Fixed per-angle contrast values, shape (n_phi,). Estimated from quantile analysis. Required for constant mode.

• fixed_offset (Array | None) – Fixed per-angle offset values, shape (n_phi,). Estimated from quantile analysis. Required for constant mode.

Return type:

None

homodyne.optimization.cmc.model.xpcs_model_averaged(data, t1, t2, phi_unique, phi_indices, q, L, dt, analysis_mode, parameter_space, n_phi, time_grid=None, noise_scale=0.1, fixed_contrast=None, fixed_offset=None, nlsq_prior_config=None, num_shards=1, shard_grid=None, **kwargs)

NumPyro model with SAMPLED averaged per-angle scaling (auto mode).

This model samples a SINGLE contrast and SINGLE offset value, then broadcasts them to all phi angles. This matches NLSQ’s auto/constant mode behavior where the averaged scaling parameters are optimized (not fixed).

Parameter count comparison (laminar_flow, n_phi=23): - individual mode: 54 params (46 per-angle + 7 physical + 1 sigma) - auto mode (this): 10 params (2 averaged scaling + 7 physical + 1 sigma) - constant mode: 8 params (7 physical + 1 sigma, scaling FIXED)

Parameters:

• data (Array) – Observed C2 correlation data, shape (n_total,).

• t1 (Array) – Time coordinates, shape (n_total,).

• t2 (Array) – Time coordinates, shape (n_total,).

• phi_unique (Array) – Unique phi angles, shape (n_phi,).

• phi_indices (Array) – Index into per-angle arrays for each point, shape (n_total,).

• q (float) – Wavevector magnitude.

• L (float) – Stator-rotor gap length (nm).

• dt (float) – Time step.

• analysis_mode (str) – Analysis mode: “static” or “laminar_flow”.

• parameter_space (ParameterSpace) – Parameter space with bounds and priors.

• n_phi (int) – Number of unique phi angles.

• noise_scale (float) – Initial estimate of observation noise.

• fixed_contrast (Array | None) – Ignored in this model. Present for API compatibility.

• fixed_offset (Array | None) – Ignored in this model. Present for API compatibility.

Return type:

None

homodyne.optimization.cmc.model.xpcs_model_constant_averaged(data, t1, t2, phi_unique, phi_indices, q, L, dt, analysis_mode, parameter_space, n_phi, time_grid=None, noise_scale=0.1, fixed_contrast=None, fixed_offset=None, nlsq_prior_config=None, num_shards=1, shard_grid=None, **kwargs)

NumPyro model with FIXED averaged per-angle scaling (NLSQ parity mode).

This model uses FIXED contrast/offset values that are the AVERAGE of per-angle estimates. These values are NOT sampled, providing exact parity with NLSQ’s “auto” mode behavior.

CRITICAL (Jan 2026): This mode fixes the parameter shift issue where CMC’s “auto” mode (xpcs_model_averaged) samples contrast/offset, introducing extra uncertainty that biases physical parameters. By using FIXED averaged values, the physical parameter posteriors should match NLSQ estimates.

Parameter count comparison (laminar_flow): - individual mode: 54 params (46 per-angle + 7 physical + 1 sigma) - auto mode (xpcs_model_averaged): 10 params (2 sampled scaling + 7 physical + 1 sigma) - constant mode (xpcs_model_constant): 8 params (7 physical + 1 sigma, per-angle fixed) - constant_averaged mode (this): 8 params (7 physical + 1 sigma, averaged fixed)

Parameters:

• data (Array) – Observed C2 correlation data, shape (n_total,).

• t1 (Array) – Time coordinates, shape (n_total,).

• t2 (Array) – Time coordinates, shape (n_total,).

• phi_unique (Array) – Unique phi angles, shape (n_phi,).

• phi_indices (Array) – Index into per-angle arrays for each point, shape (n_total,).

• q (float) – Wavevector magnitude.

• L (float) – Stator-rotor gap length (nm).

• dt (float) – Time step.

• analysis_mode (str) – Analysis mode: “static” or “laminar_flow”.

• parameter_space (ParameterSpace) – Parameter space with bounds and priors.

• n_phi (int) – Number of unique phi angles.

• noise_scale (float) – Initial estimate of observation noise.

• fixed_contrast (Array | None) – Fixed per-angle contrast values, shape (n_phi,). Will be averaged.

• fixed_offset (Array | None) – Fixed per-angle offset values, shape (n_phi,). Will be averaged.

Return type:

None

homodyne.optimization.cmc.model.xpcs_model_reparameterized(data, t1, t2, phi_unique, phi_indices, q, L, dt, analysis_mode, parameter_space, n_phi, time_grid=None, noise_scale=0.1, fixed_contrast=None, fixed_offset=None, reparam_config=None, nlsq_prior_config=None, num_shards=1, t_ref=1.0, shard_grid=None, **kwargs)

NumPyro model with reference-time reparameterized sampling space.

This model transforms correlated parameters to orthogonal sampling space: - D0, alpha → log_D_ref, alpha where D_ref = D0 * t_ref^alpha (decorrelates) - D_offset → D_offset_ratio = D_offset / D_ref (linear, handles negative D_offset) - gamma_dot_t0, beta → log_gamma_ref, beta where gamma_ref = gamma_dot_t0 * t_ref^beta

The original physics parameters (D0, D_offset, gamma_dot_t0) are computed as deterministic transforms and included in the trace for output.

D_offset_ratio uses a TruncatedNormal prior (low=-1+ε), supporting negative D_offset for jammed/arrested systems while enforcing D_ref + D_offset > 0 at t_ref. Inverse: D_offset = D_ref * ratio.

Parameters:

• reparam_config (ReparamConfig | None) – Reparameterization configuration. If None, uses defaults.

• nlsq_prior_config (dict | None) – NLSQ-informed prior configuration with keys: - “values”: dict of NLSQ parameter estimates - “uncertainties”: dict of NLSQ standard errors - “width_factor”: prior width multiplier - “reparam_values”: dict of reparameterized NLSQ values (log_D_ref, etc.) - “reparam_uncertainties”: dict of reparameterized uncertainties

• t_ref (float) – Reference time for reparameterization (default: 1.0).

• xpcs_model_averaged] ([Other parameters same as)

Return type:

None

homodyne.optimization.cmc.model.get_xpcs_model(per_angle_mode='individual', use_reparameterization=False)

Get the appropriate NumPyro model function for the given per-angle mode.

Parameters:

• per_angle_mode (str) – Per-angle scaling mode: “individual”, “auto”, “constant”, or “constant_averaged”.

• use_reparameterization (bool) – If True and per_angle_mode is “auto”, use reparameterized model for better sampling of correlated parameters (D_total instead of D0/D_offset, log_gamma_dot_t0 instead of gamma_dot_t0).

Returns:

NumPyro model function.

Return type:

callable

Notes

Mode semantics (same as NLSQ):

• individual: Uses xpcs_model_scaled which samples per-angle contrast/offset (n_phi*2 + 7 physical + 1 sigma params for laminar_flow).

• auto: Uses xpcs_model_averaged which samples SINGLE averaged contrast/offset (2 averaged + 7 physical + 1 sigma = 10 params for laminar_flow). If use_reparameterization=True, uses xpcs_model_reparameterized instead.

• constant: Uses xpcs_model_constant which requires fixed_contrast/fixed_offset arrays (NOT sampled, 7 physical + 1 sigma = 8 params for laminar_flow).

• constant_averaged: Uses xpcs_model_constant_averaged with FIXED averaged scaling (NOT sampled, 7 physical + 1 sigma = 8 params). Provides exact NLSQ parity.

CMC Per-Angle Modes

CMC supports three per-angle modes that control how contrast/offset parameters are handled during MCMC sampling. This matches the NLSQ anti-degeneracy system for consistent behavior across optimization backends.

CMC Per-Angle Mode Comparison

Mode	Sampled Params	Per-Angle Handling	Use Case
auto (default)	8 (7 physical + σ)	Quantile → average → broadcast	Default for n_phi ≥ 3 (NLSQ parity)
constant	8 (7 physical + σ)	Quantile → use directly (fixed)	Different fixed value per angle
individual	8 + 2×n_phi	All sampled independently	Full flexibility (n_phi < 3)

Auto Mode (Default):

When per_angle_mode: "auto" and n_phi ≥ 3 (configurable via constant_scaling_threshold):

1. Estimates per-angle contrast/offset from data using quantile analysis

2. Averages the per-angle estimates to single values

3. Broadcasts the averaged values to all angles (same fixed value for all)

4. Only samples 8 parameters: 7 physical + 1 sigma

This provides NLSQ parity—CMC auto mode matches NLSQ constant mode behavior.

Constant Mode:

When per_angle_mode: "constant":

1. Estimates per-angle contrast/offset from data using quantile analysis

2. Uses the per-angle estimates directly (different fixed value per angle)

3. Only samples 8 parameters: 7 physical + 1 sigma

Both auto (n_phi ≥ 3) and constant modes use fixed scaling arrays passed to the model function, reducing degeneracy risk by not sampling per-angle parameters.

Individual Mode:

When per_angle_mode: "individual" or auto with n_phi < 3:

1. Samples contrast and offset for each phi angle independently

2. Total sampled parameters: 8 + 2×n_phi

3. Full flexibility but higher degeneracy risk for large n_phi

Model Selection:

from homodyne.optimization.cmc.model import get_xpcs_model

# Get appropriate model function
model = get_xpcs_model("constant")  # Returns xpcs_model_constant
model = get_xpcs_model("individual")  # Returns xpcs_model_scaled
model = get_xpcs_model()  # Default: xpcs_model_scaled

Configuration:

optimization:
  cmc:
    per_angle_mode: "auto"           # "auto", "constant", "individual"
    constant_scaling_threshold: 3    # Threshold for auto mode

Key Functions

homodyne.optimization.cmc.model.xpcs_model_constant	NumPyro model with FIXED per-angle scaling (anti-degeneracy constant mode).
homodyne.optimization.cmc.model.xpcs_model_scaled	NumPyro model with non-centered parameterization for gradient balancing.
homodyne.optimization.cmc.model.get_xpcs_model	Get the appropriate NumPyro model function for the given per-angle mode.
homodyne.optimization.cmc.model.get_model_param_count	Get total number of sampled parameters.
homodyne.optimization.cmc.priors.get_param_names_in_order	Get parameter names in NumPyro sampling order.
homodyne.optimization.cmc.priors.build_init_values_dict	Build complete initial values dictionary in sampling order.
homodyne.optimization.cmc.config.CMCConfig.get_effective_per_angle_mode	Determine effective per-angle mode based on configuration and data.

CMC Convergence and Precision Fixes

This section documents comprehensive fixes for CMC failures on multi-angle datasets, addressing 94% shard timeout rates, 28.4% divergence rates, and 33-43x uncertainty inflation observed in 3-angle laminar_flow analysis.

Angle-Aware Shard Sizing

The _resolve_max_points_per_shard() function now accepts an n_phi parameter that scales shard sizes inversely with angle count:

Angle-Aware Shard Scaling

n_phi	Scale Factor	Rationale
≤ 3	30%	Few-angle data has more complex per-shard posteriors
4-5	50%	Moderate angle count
6-10	70%	Good angle coverage per shard
> 10	100%	Full capacity

Example: For 3-angle laminar_flow with base 20K shard size, effective size = 6K points.

Angle-Balanced Sharding

New shard_data_angle_balanced() function ensures proportional angle coverage per shard:

from homodyne.optimization.cmc.data_prep import shard_data_angle_balanced

shards = shard_data_angle_balanced(
    prepared,
    num_shards=None,           # Auto-calculate
    max_points_per_shard=6000, # Angle-aware size
    min_angle_coverage=0.8,    # 80% minimum coverage
    seed=42,
)

Key features:

• Samples proportionally from each angle group

• Logs coverage statistics per shard

• Falls back to random sharding if angle-balanced impossible

NLSQ Warm-Start Priors

New functions in homodyne.optimization.cmc.priors for NLSQ-informed prior construction:

from homodyne.optimization.cmc.priors import (
    build_nlsq_informed_prior,
    build_nlsq_informed_priors,
    extract_nlsq_values_for_cmc,
)

# Extract NLSQ values from various result formats
nlsq_values = extract_nlsq_values_for_cmc(nlsq_result)

# Build informative prior for single parameter
prior = build_nlsq_informed_prior(
    param_name="D0",
    nlsq_value=1234.5,
    nlsq_std=45.6,
    bounds=(100, 10000),
    width_factor=3.0,  # 3σ width
)

# Build priors for all physical parameters
priors = build_nlsq_informed_priors(nlsq_values, nlsq_stds, bounds, analysis_mode)

Usage in fit_mcmc_jax:

from homodyne.optimization.cmc import fit_mcmc_jax
from homodyne.optimization.nlsq import fit_nlsq_jax

# Step 1: Run NLSQ
nlsq_result = fit_nlsq_jax(data, config)

# Step 2: Run CMC with NLSQ warm-start
cmc_result = fit_mcmc_jax(data, config, nlsq_result=nlsq_result)

Constant-Averaged Per-Angle Mode

New xpcs_model_constant_averaged() model for exact NLSQ “auto” mode parity:

• Uses FIXED averaged contrast/offset (not sampled)

• 8 parameters (7 physical + sigma) instead of 10

• Matches NLSQ constant mode averaging behavior

optimization:
  cmc:
    per_angle_mode: "constant_averaged"  # Match NLSQ "auto"

Early Abort Mechanism

The multiprocessing backend now tracks failure categories and aborts early:

Failure Categories

Category	Description
timeout	Shard exceeded per_shard_timeout
heartbeat_timeout	Worker stopped responding
crash	Worker process crashed
numerical	NaN/Inf in posterior samples
convergence	High R-hat or low ESS

Abort condition: If >50% of first 10 shards fail, the run aborts immediately.

NUTS Convergence Improvements

For laminar_flow mode:

• target_accept_prob automatically elevated to 0.9 (from default 0.85)

• Divergence rate checking with severity levels:

  • >30%: CRITICAL (logged as error, run continues)

  • >10%: WARNING

  • >5%: ELEVATED (info)

Precision Diagnostics

New functions in homodyne.optimization.cmc.diagnostics:

from homodyne.optimization.cmc.diagnostics import (
    compute_posterior_contraction,
    compute_nlsq_comparison_metrics,
    compute_precision_analysis,
    log_precision_analysis,
)

# Posterior Contraction Ratio: PCR = 1 - (posterior_std / prior_std)
pcr = compute_posterior_contraction(posterior_std=10.0, prior_std=100.0)
# pcr = 0.9 (90% contraction = informative data)

# Compare CMC to NLSQ
metrics = compute_nlsq_comparison_metrics(
    cmc_mean=1234.5,
    cmc_std=45.6,
    nlsq_value=1250.0,
    nlsq_std=50.0,
)
# Returns: z_score, uncertainty_ratio, overlap

Configuration Reference

optimization:
  cmc:
    sharding:
      max_points_per_shard: "auto"  # Angle-aware scaling
      strategy: "angle_balanced"    # Ensure coverage per shard
      min_angle_coverage: 0.8       # 80% of angles per shard
    sampler:
      target_accept_prob: 0.9       # Higher for laminar_flow
    execution:
      per_shard_timeout: 3600       # 1 hour (down from 2)
      early_abort_threshold: 0.5    # Abort if >50% of first 10 fail
    per_angle_mode: "constant_averaged"  # Match NLSQ "auto"

Shared Scaling Utilities

CMC uses shared utilities from homodyne.core.scaling_utils for quantile-based contrast/offset estimation:

Shared scaling utilities for per-angle contrast/offset estimation.

This module provides unified quantile-based estimation of contrast and offset parameters that can be used by both NLSQ and CMC optimization backends.

The physics basis:

C2 = contrast × g1² + offset

• At large time lags, g1² → 0, so C2 → offset (the “floor”)

• At small time lags, g1² ≈ 1, so C2 ≈ contrast + offset (the “ceiling”)

Version: 2.18.0

homodyne.core.scaling_utils.estimate_contrast_offset_from_quantiles(c2_data, delta_t, contrast_bounds=(0.0, 1.0), offset_bounds=(0.5, 1.5), lag_floor_quantile=0.80, lag_ceiling_quantile=0.20, value_quantile_low=0.10, value_quantile_high=0.90)

Estimate contrast and offset from C2 data using physics-informed quantile analysis.

Uses the correlation decay structure: C2 = contrast × g1² + offset - At large time lags, g1² → 0, so C2 → offset (the “floor”) - At small time lags, g1² ≈ 1, so C2 ≈ contrast + offset (the “ceiling”)

Parameters:

• c2_data (ndarray) – C2 correlation values (1D array).

• delta_t (ndarray) – Time lag values abs(t1 - t2) (same shape as c2_data).

• contrast_bounds (tuple[float, float]) – Valid bounds for contrast parameter.

• offset_bounds (tuple[float, float]) – Valid bounds for offset parameter.

• lag_floor_quantile (float) – Quantile threshold for “large lag” region (default: 0.80 = top 20% of lags).

• lag_ceiling_quantile (float) – Quantile threshold for “small lag” region (default: 0.20 = bottom 20% of lags).

• value_quantile_low (float) – Quantile for robust floor estimation (default: 0.10).

• value_quantile_high (float) – Quantile for robust ceiling estimation (default: 0.90).

Returns:

(contrast_est, offset_est) - Estimated values clipped to bounds.

Return type:

tuple[float, float]

Notes

The estimation is robust to outliers by using quantiles instead of min/max. The lag-based segmentation ensures we’re sampling from the appropriate regions of the correlation decay curve.

homodyne.core.scaling_utils.estimate_per_angle_scaling(c2_data, t1, t2, phi_indices, n_phi, contrast_bounds, offset_bounds, log=None)

Estimate contrast and offset initial values for each phi angle.

This is the unified implementation used by both NLSQ and CMC backends.

Optimization (v2.9.1): Uses vectorized grouped operations instead of sequential loop over angles. Provides 3-5x speedup for typical datasets with 20+ phi angles.

Parameters:

• c2_data (ndarray) – Pooled C2 correlation values.

• t1 (ndarray) – Pooled first time coordinates.

• t2 (ndarray) – Pooled second time coordinates.

• phi_indices (ndarray) – Index mapping each data point to its phi angle (0 to n_phi-1).

• n_phi (int) – Number of unique phi angles.

• contrast_bounds (tuple[float, float]) – Valid bounds for contrast.

• offset_bounds (tuple[float, float]) – Valid bounds for offset.

• log (Logger | LoggerAdapter[Logger] | None) – Logger for diagnostic messages.

Returns:

Dictionary with keys ‘contrast_0’, ‘offset_0’, ‘contrast_1’, ‘offset_1’, etc.

Return type:

dict[str, float]

homodyne.core.scaling_utils.compute_averaged_scaling(c2_data, t1, t2, phi_indices, n_phi, contrast_bounds, offset_bounds, log=None)

Compute averaged contrast and offset for constant mode.

This function estimates per-angle contrast/offset using quantile analysis, then averages them to produce single values for constant mode optimization.

Parameters:

• c2_data (ndarray) – Pooled C2 correlation values.

• t1 (ndarray) – Pooled first time coordinates.

• t2 (ndarray) – Pooled second time coordinates.

• phi_indices (ndarray) – Index mapping each data point to its phi angle (0 to n_phi-1).

• n_phi (int) – Number of unique phi angles.

• contrast_bounds (tuple[float, float]) – Valid bounds for contrast.

• offset_bounds (tuple[float, float]) – Valid bounds for offset.

• log (Logger | LoggerAdapter[Logger] | None) – Logger for diagnostic messages.

Returns:

(contrast_avg, offset_avg, contrast_per_angle, offset_per_angle) - contrast_avg: Averaged contrast for constant mode - offset_avg: Averaged offset for constant mode - contrast_per_angle: Per-angle estimates (for diagnostics) - offset_per_angle: Per-angle estimates (for diagnostics)

Return type:

tuple[float, float, ndarray, ndarray]

Priors

Physics-informed prior distributions.

Prior distribution builders for CMC analysis.

This module provides utilities for building NumPyro prior distributions from the ParameterSpace configuration.

homodyne.optimization.cmc.priors.estimate_contrast_offset_from_data(c2_data, t1, t2, contrast_bounds=(0.0, 1.0), offset_bounds=(0.5, 1.5), lag_floor_quantile=0.80, lag_ceiling_quantile=0.20, value_quantile_low=0.10, value_quantile_high=0.90)

Estimate contrast and offset from C2 data using physics-informed quantile analysis.

Uses the correlation decay structure: C2 = contrast × g1² + offset - At large time lags, g1² → 0, so C2 → offset (the “floor”) - At small time lags, g1² ≈ 1, so C2 ≈ contrast + offset (the “ceiling”)

Parameters:

• c2_data (ndarray) – C2 correlation values (1D array).

• t1 (ndarray) – First time coordinate array (same shape as c2_data).

• t2 (ndarray) – Second time coordinate array (same shape as c2_data).

• contrast_bounds (tuple[float, float]) – Valid bounds for contrast parameter.

• offset_bounds (tuple[float, float]) – Valid bounds for offset parameter.

• lag_floor_quantile (float) – Quantile threshold for “large lag” region (default: 0.80 = top 20% of lags).

• lag_ceiling_quantile (float) – Quantile threshold for “small lag” region (default: 0.20 = bottom 20% of lags).

• value_quantile_low (float) – Quantile for robust floor estimation (default: 0.10).

• value_quantile_high (float) – Quantile for robust ceiling estimation (default: 0.90).

Returns:

(contrast_est, offset_est) - Estimated values clipped to bounds.

Return type:

tuple[float, float]

Notes

The estimation is robust to outliers by using quantiles instead of min/max. The lag-based segmentation ensures we’re sampling from the appropriate regions of the correlation decay curve.

homodyne.optimization.cmc.priors.estimate_per_angle_scaling(c2_data, t1, t2, phi_indices, n_phi, contrast_bounds, offset_bounds)

Estimate contrast and offset initial values for each phi angle.

Thin wrapper that delegates to the canonical implementation in homodyne.core.scaling_utils. Kept here for backward compatibility with any internal callers within this module.

Parameters:

• c2_data (ndarray) – Pooled C2 correlation values.

• t1 (ndarray) – Pooled first time coordinates.

• t2 (ndarray) – Pooled second time coordinates.

• phi_indices (ndarray) – Index mapping each data point to its phi angle (0 to n_phi-1).

• n_phi (int) – Number of unique phi angles.

• contrast_bounds (tuple[float, float]) – Valid bounds for contrast.

• offset_bounds (tuple[float, float]) – Valid bounds for offset.

Returns:

Dictionary with keys ‘contrast_0’, ‘offset_0’, ‘contrast_1’, ‘offset_1’, etc.

Return type:

dict[str, float]

homodyne.optimization.cmc.priors.build_prior_from_spec(prior_spec)

Build NumPyro distribution from PriorDistribution specification.

Parameters:

prior_spec (PriorDistribution) – Prior specification from ParameterSpace.

Returns:

NumPyro distribution object.

Return type:

Distribution

Raises:

ValueError – If distribution type is not supported.

homodyne.optimization.cmc.priors.build_prior(param_name, parameter_space)

Build NumPyro prior distribution for a parameter.

Parameters:

• param_name (str) – Parameter name (e.g., “D0”, “alpha”, “contrast”, “contrast_0”).

• parameter_space (ParameterSpace) – Parameter space with bounds and priors.

Returns:

NumPyro distribution for sampling.

Return type:

Distribution

homodyne.optimization.cmc.priors.get_init_value(param_name, initial_values, parameter_space)

Get initial value for a parameter.

Priority: 1. Value from initial_values dict if provided (exact match) 2. Value from initial_values dict for base param (e.g., ‘contrast’ for ‘contrast_0’) 3. Midpoint of parameter bounds as fallback

Parameters:

• param_name (str) – Parameter name.

• initial_values (dict[str, float] | None) – Initial values from config.

• parameter_space (ParameterSpace) – Parameter space with bounds.

Returns:

Initial value for the parameter.

Return type:

float

Notes

Per-angle parameter handling (scalar broadcast):

For per-angle parameters like ‘contrast_0’, ‘contrast_1’, etc., this function broadcasts a single scalar value to all angles. If only ‘contrast’ is provided in initial_values (not ‘contrast_0’, ‘contrast_1’, etc.), that single value is used for ALL phi angles.

To specify different initial values per angle, provide explicit keys like: {'contrast_0': 0.4, 'contrast_1': 0.5, 'contrast_2': 0.45}

The same applies to ‘offset’ parameters.

Examples

>>> # Scalar broadcast: same value for all angles
>>> initial_values = {'contrast': 0.5, 'offset': 1.0}
>>> get_init_value('contrast_0', initial_values, param_space)  # Returns 0.5
>>> get_init_value('contrast_1', initial_values, param_space)  # Returns 0.5

>>> # Explicit per-angle values
>>> initial_values = {'contrast_0': 0.4, 'contrast_1': 0.6}
>>> get_init_value('contrast_0', initial_values, param_space)  # Returns 0.4
>>> get_init_value('contrast_1', initial_values, param_space)  # Returns 0.6

homodyne.optimization.cmc.priors.validate_initial_value_bounds(param_name, value, parameter_space)

Validate and optionally clip initial value to parameter bounds.

Parameters:

• param_name (str) – Parameter name.

• value (float) – Initial value to validate.

• parameter_space (ParameterSpace) – Parameter space with bounds.

Returns:

(validated_value, was_clipped) - The value (clipped if needed) and whether clipping occurred.

Return type:

tuple[float, bool]

homodyne.optimization.cmc.priors.build_init_values_dict(n_phi, analysis_mode, initial_values, parameter_space, *, c2_data=None, t1=None, t2=None, phi_indices=None, per_angle_mode='individual')

Build complete initial values dictionary in sampling order.

CRITICAL: Parameter order must match NumPyro model sampling order:

1. contrast_0, contrast_1, …, contrast_{n_phi-1} (individual mode) OR contrast_avg (constant mode).

2. offset_0, offset_1, …, offset_{n_phi-1} (individual mode) OR offset_avg (constant mode).

3. Physical parameters in canonical order.

Parameters:

• n_phi (int) – Number of phi angles.

• analysis_mode (str) – Analysis mode (“static” or “laminar_flow”).

• initial_values (dict[str, float] | None) – Initial values from config. Supports both scalar (broadcast) and per-angle specifications for contrast/offset. See Notes for details.

• parameter_space (ParameterSpace) – Parameter space with bounds.

• c2_data (ndarray | None) – Optional C2 correlation data for quantile-based estimation of contrast/offset.

• t1 (ndarray | None) – Optional time coordinates (required if c2_data provided).

• t2 (ndarray | None) – Optional time coordinates (required if c2_data provided).

• phi_indices (ndarray | None) – Optional phi angle indices for per-angle estimation.

• per_angle_mode (str) – Per-angle scaling mode: “individual” or “constant”.

Returns:

Initial values dictionary in sampling order.

Return type:

dict[str, float]

Notes

Per-angle scaling parameters (contrast/offset):

This function supports three modes for specifying per-angle initial values:

1. Data-driven estimation (NEW, preferred): If c2_data, t1, t2, and phi_indices are provided, and contrast/offset not in initial_values, uses physics-informed quantile analysis to estimate values from data.

2. Scalar broadcast: If initial_values contains only base names like ‘contrast’ and ‘offset’, those values are broadcast to ALL phi angles. Example: {'contrast': 0.5} → contrast_0=0.5, contrast_1=0.5, …

3. Explicit per-angle: If initial_values contains indexed names like ‘contrast_0’, ‘contrast_1’, etc., those specific values are used. Example: {'contrast_0': 0.4, 'contrast_1': 0.6}

Priority: explicit per-angle > scalar broadcast > data-driven > midpoint fallback

Bounds validation:

All initial values are validated against parameter bounds. Out-of-bounds values are clipped to bounds ± 1% margin with a warning logged.

homodyne.optimization.cmc.priors.get_param_names_in_order(n_phi, analysis_mode, per_angle_mode='individual')

Get parameter names in NumPyro sampling order.

CRITICAL: This order must match the model sampling order exactly.

Parameters:

• n_phi (int) – Number of phi angles.

• analysis_mode (str) – Analysis mode (“static” or “laminar_flow”).

• per_angle_mode (str) – Per-angle scaling mode: “individual”, “auto”, or “constant”.

Returns:

Parameter names in sampling order.

Return type:

list[str]

Notes

Mode semantics (same as NLSQ): - individual mode: Samples per-angle contrast/offset (2*n_phi params) - auto mode: Samples single averaged contrast/offset (2 params) - constant mode: NO contrast/offset sampled (fixed from quantile estimation)

homodyne.optimization.cmc.priors.validate_init_values_order(init_values, expected_names)

Validate that init values dictionary keys match expected order.

This is a defensive check to catch parameter ordering bugs early. In Python 3.7+, dict preserves insertion order, so key order matters for functions that assume positional correspondence.

Parameters:

• init_values (dict[str, float]) – Initial values dictionary.

• expected_names (list[str]) – Expected parameter names in order.

Raises:

ValueError – If parameter order doesn’t match.

Return type:

None

homodyne.optimization.cmc.priors.build_nlsq_informed_prior(param_name, nlsq_value, nlsq_std, bounds, width_factor=2.0)

Build a TruncatedNormal prior centered on NLSQ estimate.

This provides informative priors for CMC that leverage NLSQ’s point estimates. The resulting priors: 1. Center at the NLSQ estimate (faster warmup, better mixing) 2. Have width based on NLSQ uncertainty or parameter range 3. Are truncated to respect parameter bounds 4. Enable posterior contraction metrics (comparing prior vs posterior width)

Parameters:

• param_name (str) – Parameter name for logging.

• nlsq_value (float) – NLSQ point estimate (mean of the prior).

• nlsq_std (float | None) – NLSQ standard error estimate. If None, uses 10% of bounds range.

• bounds (tuple[float, float]) – Parameter bounds (low, high).

• width_factor (float) – Multiplier for NLSQ std to get prior width. Default 2.0 gives ~95% coverage assuming Gaussian posterior.

Returns:

TruncatedNormal distribution centered at nlsq_value.

Return type:

Distribution

homodyne.optimization.cmc.priors.build_nlsq_informed_priors(nlsq_result, nlsq_uncertainties, parameter_space, analysis_mode, n_phi, width_factor=2.0)

Build informative priors for all physical parameters from NLSQ results.

Parameters:

• nlsq_result (dict[str, float]) – NLSQ parameter estimates (e.g., {“D0”: 1e10, “alpha”: -0.5, …}).

• nlsq_uncertainties (dict[str, float] | None) – NLSQ standard errors for each parameter. If None, uses weak priors.

• parameter_space (ParameterSpace) – Parameter space with bounds.

• analysis_mode (str) – Analysis mode: “static” or “laminar_flow”.

• n_phi (int) – Number of phi angles (for per-angle parameters if needed).

• width_factor (float) – Width multiplier for priors. Default 2.0.

Returns:

Dictionary of informative priors keyed by parameter name.

Return type:

dict[str, Distribution]

homodyne.optimization.cmc.priors.extract_nlsq_values_for_cmc(nlsq_result)

Extract parameter values and uncertainties from an NLSQ result.

This utility handles various NLSQ result formats and extracts the information needed for CMC warm-start priors.

Parameters:

nlsq_result (dict | Any) – NLSQ result, either: - OptimizationResult dataclass with parameters/uncertainties arrays - dict with “params”/”parameters”/”best_params” keys - dict with flat structure (parameter names as keys)

Returns:

Tuple of (parameter_values, uncertainties). uncertainties may be None if not available.

Return type:

tuple[dict[str, float], dict[str, float] | None]

Prior Specifications

Static Mode (3 physical parameters):

• D0: LogNormal(log(1000), 1.5)

• alpha: Uniform(0.0, 2.0)

• D_offset: TruncatedNormal(0, 100, low=0)

Laminar Flow Mode (+4 shear parameters):

• gamma_dot_t0: LogNormal(log(100), 1.5)

• beta: Uniform(-2.0, 2.0)

• gamma_dot_t_offset: TruncatedNormal(0, 100, low=0)

• phi0: Uniform(0, 2π)

Per-Angle Scaling (mandatory):

• contrast_i: TruncatedNormal(0.5, 0.3, low=0.1, high=2.0) for each angle i

• offset_i: TruncatedNormal(1.0, 0.2, low=0.5, high=1.5) for each angle i

Data Preparation

Data preparation and validation for CMC analysis.

This module provides utilities for validating and preparing pooled XPCS data for MCMC sampling.

class homodyne.optimization.cmc.data_prep.PreparedData

Bases: object

Validated and prepared data for MCMC sampling.

data

Pooled C2 correlation data, shape (n_total,).

Type:

np.ndarray

t1

Pooled time coordinates t1, shape (n_total,).

Type:

np.ndarray

t2

Pooled time coordinates t2, shape (n_total,).

Type:

np.ndarray

phi

Pooled phi angles, shape (n_total,).

Type:

np.ndarray

phi_unique

Unique phi angles, shape (n_phi,).

Type:

np.ndarray

phi_indices

Index of phi_unique for each data point, shape (n_total,).

Type:

np.ndarray

n_total

Total number of data points.

Type:

int

n_phi

Number of unique phi angles.

Type:

int

noise_scale

Estimated observation noise scale.

Type:

float

data: ndarray

t1: ndarray

t2: ndarray

phi: ndarray

phi_unique: ndarray

phi_indices: ndarray

n_total: int

n_phi: int

noise_scale: float

__init__(data, t1, t2, phi, phi_unique, phi_indices, n_total, n_phi, noise_scale)

homodyne.optimization.cmc.data_prep.validate_pooled_data(data, t1, t2, phi)

Validate that pooled data arrays are consistent.

Parameters:

• data (ndarray) – Pooled C2 correlation data.

• t1 (ndarray) – Pooled time coordinates t1.

• t2 (ndarray) – Pooled time coordinates t2.

• phi (ndarray) – Pooled phi angles.

Raises:

ValueError – If arrays have inconsistent shapes or contain invalid values.

Return type:

None

homodyne.optimization.cmc.data_prep.extract_phi_info(phi)

Extract unique phi angles and compute index mapping.

Parameters:

phi (ndarray) – Pooled phi angles, shape (n_total,).

Returns:

• phi_unique: Unique phi values sorted, shape (n_phi,)

• phi_indices: Index into phi_unique for each point, shape (n_total,)

Return type:

tuple[ndarray, ndarray]

homodyne.optimization.cmc.data_prep.estimate_noise_scale(data)

Estimate observation noise scale from data.

Uses robust MAD (Median Absolute Deviation) estimator scaled to approximate standard deviation for Gaussian noise.

Parameters:

data (ndarray) – Observed data values.

Returns:

Estimated noise scale (standard deviation).

Return type:

float

homodyne.optimization.cmc.data_prep.compute_data_statistics(data)

Compute summary statistics for data.

Parameters:

data (ndarray) – Data array.

Returns:

Statistics including min, max, mean, std, median.

Return type:

dict[str, float]

homodyne.optimization.cmc.data_prep.prepare_mcmc_data(data, t1, t2, phi, filter_diagonal=True)

Prepare and validate data for MCMC sampling.

Parameters:

• data (ndarray) – Pooled C2 correlation data, shape (n_total,).

• t1 (ndarray) – Pooled time coordinates t1, shape (n_total,).

• t2 (ndarray) – Pooled time coordinates t2, shape (n_total,).

• phi (ndarray) – Pooled phi angles, shape (n_total,).

• filter_diagonal (bool) – If True, exclude diagonal points (t1 == t2) from the dataset. Diagonal points represent autocorrelation peaks that are corrected at load time but should not contribute to the likelihood. Added in v2.14.2 for consistency with NLSQ diagonal handling.

Returns:

Validated and prepared data object.

Return type:

PreparedData

Raises:

ValueError – If data validation fails.

homodyne.optimization.cmc.data_prep.shard_data_stratified(prepared, num_shards=None, max_points_per_shard=None, max_shards_per_angle=100, seed=42)

Shard data by phi angle (stratified sharding).

Each shard contains data for one phi angle. If a single angle has more data points than max_points_per_shard, multiple shards are created for that angle by splitting the data randomly.

When the number of required shards exceeds max_shards_per_angle, shard size increases to fit all data (no subsampling).

Parameters:

• prepared (PreparedData) – Prepared data object.

• num_shards (int | None) – Desired total shard count. When provided, it forces a target shard size; max_points_per_shard is derived if not set.

• max_points_per_shard (int | None) – Maximum points per shard. If an angle exceeds this, multiple shards are created for that angle. If None, uses one shard per angle unless num_shards is provided (then it is derived). Recommended: 25000-100000 for NUTS.

• max_shards_per_angle (int) – Maximum shards to create per angle. If more would be needed, shard size increases to fit all data. Default: 100.

• seed (int) – Random seed for reproducible splitting.

Returns:

List of shard data objects.

Return type:

list[PreparedData]

homodyne.optimization.cmc.data_prep.shard_data_random(prepared, num_shards=None, max_points_per_shard=None, max_shards=100, seed=42)

Shard data randomly into approximately equal parts.

This is used when there’s only one phi angle but the dataset is too large for efficient NUTS sampling. Each shard gets a random subset of the data. ALL data is used (no subsampling).

Parameters:

• prepared (PreparedData) – Prepared data object.

• num_shards (int | None) – Number of shards to create. If None, calculated from data size and max_points_per_shard.

• max_points_per_shard (int | None) – Target points per shard. Used to calculate num_shards if not provided. If num_shards would exceed max_shards, shard size increases to fit all data. Recommended: 25000-100000 for NUTS.

• max_shards (int) – Maximum number of shards. Default: 100.

• seed (int) – Random seed for reproducible shuffling.

Returns:

List of shard data objects.

Return type:

list[PreparedData]

homodyne.optimization.cmc.data_prep.shard_data_angle_balanced(prepared, num_shards=None, max_points_per_shard=None, max_shards=500, min_angle_coverage=0.8, seed=42)

Shard data with balanced angle coverage per shard.

This is the preferred sharding method for multi-angle datasets (n_phi > 1) when using random/mixed sharding. Unlike pure random sharding, this method ensures each shard has representative coverage from each phi angle.

CRITICAL (Jan 2026): Prevents heterogeneous posteriors that cause high CV across shards. The D_offset CV=1.58 failure case was caused by pure random sharding creating shards with uneven angle coverage.

Algorithm: 1. Shuffle data within each angle independently 2. For each shard, sample proportionally from each angle 3. Verify angle coverage meets minimum threshold 4. Log coverage statistics for diagnostics

Parameters:

• prepared (PreparedData) – Prepared data object with multi-angle data.

• num_shards (int | None) – Number of shards to create. If None, calculated from data size and max_points_per_shard.

• max_points_per_shard (int | None) – Target points per shard. Used to calculate num_shards if not provided. Recommended: 3000-10000 for laminar_flow with few angles.

• max_shards (int) – Maximum number of shards. Default: 500 (higher than random to allow smaller shards for multi-angle data).

• min_angle_coverage (float) – Minimum fraction of angles that must be present in each shard. Default: 0.8 (80% of angles). Shards below this threshold are logged.

• seed (int) – Random seed for reproducible sampling.

Returns:

List of shard data objects, each with balanced angle coverage.

Return type:

list[PreparedData]

Notes

• ALL data is used (no subsampling)

• Each shard aims to have the same proportion of each angle as the full dataset

• The last shard may have slightly different proportions to include all data

homodyne.optimization.cmc.data_prep.create_xdata_dict(prepared, q, L, dt, analysis_mode)

Create xdata dictionary for physics model.

Parameters:

• prepared (PreparedData) – Prepared data object.

• q (float) – Wavevector magnitude.

• L (float) – Stator-rotor gap length.

• dt (float) – Time step.

• analysis_mode (str) – Analysis mode (“static” or “laminar_flow”).

Returns:

Dictionary of model inputs.

Return type:

dict[str, Any]

Sampler

NUTS sampler interface with warmup and sampling phases.

NUTS sampler wrapper for CMC analysis.

This module provides utilities for running NumPyro NUTS sampling with proper initialization and progress tracking.

class homodyne.optimization.cmc.sampler.SamplingStats

Bases: object

Statistics from MCMC sampling.

warmup_time

Time spent in warmup phase (seconds).

Type:

float

sampling_time

Time spent in sampling phase (seconds).

Type:

float

total_time

Total sampling time (seconds).

Type:

float

num_divergent

Number of divergent transitions.

Type:

int

accept_prob

Mean acceptance probability.

Type:

float

step_size

Final step size.

Type:

float

step_size_min

Minimum adapted step size across chains (if available).

Type:

float

step_size_max

Maximum adapted step size across chains (if available).

Type:

float

inverse_mass_matrix_summary

Compact summary of the adapted inverse mass matrix (if available).

Type:

str | None

tree_depth

Mean tree depth.

Type:

float

warmup_time: float = 0.0

sampling_time: float = 0.0

total_time: float = 0.0

num_divergent: int = 0

accept_prob: float = 0.0

step_size: float = 0.0

step_size_min: float | None = None

step_size_max: float | None = None

inverse_mass_matrix_summary: str | None = None

tree_depth: float = 0.0

plan: SamplingPlan | None = None

__init__(warmup_time=0.0, sampling_time=0.0, total_time=0.0, num_divergent=0, accept_prob=0.0, step_size=0.0, step_size_min=None, step_size_max=None, inverse_mass_matrix_summary=None, tree_depth=0.0, plan=None)

class homodyne.optimization.cmc.sampler.SamplingPlan

Bases: object

Adapted MCMC sampling counts for a single shard.

Captures the actual warmup/sample counts after adaptive scaling, which may differ from CMCConfig defaults for small shards.

Use SamplingPlan.from_config() instead of accessing config.num_warmup / config.num_samples in hot paths.

n_warmup: int

n_samples: int

n_chains: int

shard_size: int

n_params: int

was_adapted: bool

classmethod from_config(config, shard_size, n_params)

Return type:

SamplingPlan

property total_samples: int

__init__(n_warmup, n_samples, n_chains, shard_size, n_params, was_adapted)

class homodyne.optimization.cmc.sampler.MCMCSamples

Bases: object

Container for MCMC samples.

samples

Parameter samples, shape (n_chains, n_samples) per parameter.

Type:

dict[str, np.ndarray]

param_names

Parameter names in sampling order.

Type:

list[str]

n_chains

Number of chains.

Type:

int

n_samples

Number of samples per chain.

Type:

int

extra_fields

Additional MCMC info (divergences, energy, etc.).

Type:

dict[str, Any]

num_shards

Number of shards combined (1 for single shard, >1 for CMC). Used for correct divergence rate calculation in CMC.

Type:

int

samples: dict[str, ndarray]

param_names: list[str]

n_chains: int

n_samples: int

extra_fields: dict[str, Any]

num_shards: int = 1

shard_adapted_n_warmup: int | None = None

bimodal_consensus: Any = None

__init__(samples, param_names, n_chains, n_samples, extra_fields=<factory>, num_shards=1, shard_adapted_n_warmup=None, bimodal_consensus=None)

homodyne.optimization.cmc.sampler.create_init_strategy(initial_values, param_names, use_init_to_value=True, z_space_values=None)

Create initialization strategy for NUTS.

Parameters:

• initial_values (dict[str, float] | None) – Initial values from config (original space).

• param_names (list[str]) – Expected parameter names in order.

• use_init_to_value (bool) – If True, use init_to_value when values provided.

• z_space_values (dict[str, float] | None) – Initial values in z-space (for scaled model). If provided, these are used directly as {name}_z values.

Returns:

NumPyro initialization function.

Return type:

Callable

homodyne.optimization.cmc.sampler.run_nuts_sampling(model, model_kwargs, config, initial_values, parameter_space, n_phi, analysis_mode, rng_key=None, progress_bar=True, per_angle_mode='individual')

Run NUTS sampling with configuration.

Parameters:

• model (Callable) – NumPyro model function.

• model_kwargs (dict[str, Any]) – Keyword arguments to pass to model.

• config (CMCConfig) – CMC configuration.

• initial_values (dict[str, float] | None) – Initial parameter values from config.

• parameter_space (ParameterSpace) – Parameter space for building init values.

• n_phi (int) – Number of phi angles.

• analysis_mode (str) – Analysis mode.

• rng_key (jax.random.PRNGKey | None) – Random key. If None, creates from seed.

• progress_bar (bool) – Whether to show progress bar.

• per_angle_mode (str) – Per-angle scaling mode: “individual”, “auto”, “constant”, or “constant_averaged”. Controls which parameters are sampled vs fixed.

Returns:

Samples and timing statistics.

Return type:

tuple[MCMCSamples, SamplingStats]

homodyne.optimization.cmc.sampler.run_nuts_with_retry(model, model_kwargs, config, initial_values, parameter_space, n_phi, analysis_mode, max_retries=3, rng_key=None, per_angle_mode='individual')

Run NUTS sampling with automatic retry on failure.

Parameters:

• model (Callable) – NumPyro model function.

• model_kwargs (dict[str, Any]) – Model arguments.

• config (CMCConfig) – Configuration.

• initial_values (dict[str, float] | None) – Initial values.

• parameter_space (ParameterSpace) – Parameter space.

• n_phi (int) – Number of phi angles.

• analysis_mode (str) – Analysis mode.

• max_retries (int) – Maximum number of retry attempts.

• rng_key (jax.random.PRNGKey | None) – Random key.

Returns:

Samples and statistics.

Return type:

tuple[MCMCSamples, SamplingStats]

Raises:

RuntimeError – If all retries fail.

Parameter Scaling (Gradient Balancing)

Parameter Scaling for MCMC Gradient Balancing.

This module implements non-centered reparameterization to balance gradient scales across parameters with vastly different magnitudes.

The Problem:

In the CMC model, parameters span many orders of magnitude: - D0: ~10^4 (diffusion coefficient) - alpha: ~10^0 (exponent) - gamma_dot_t0: ~10^-3 (shear rate) - contrast: ~10^-1 (optical scaling)

When NUTS samples these parameters directly, gradients are dominated by large-scale parameters (D0), causing the sampler to effectively ignore small-scale parameters. This leads to 0% acceptance rate.

The Solution:

Non-centered reparameterization transforms each parameter to unit scale:

P_z ~ Normal(0, 1) # Sample in normalized space P = center + scale × P_z # Transform to original space P = smooth_bound(P, low, high) # Smoothly enforce bounds

Where: - center = (low + high) / 2 or prior_mu - scale = (high - low) / 4 or prior_sigma

This ensures ALL gradients have similar magnitude, enabling balanced MCMC exploration.

CRITICAL - Lessons Learned (Dec 2025):

Hard clipping (jnp.clip) introduces non-smooth behavior at the bounds. In practice this can lead to poor HMC/NUTS adaptation (especially when chains push against bounds during warmup), including near-zero acceptance.

To avoid this, Homodyne uses a smooth bounded transform based on tanh:

smooth_bound(x; low, high) = mid + half * tanh((x - mid) / half)

This maps ℝ → (low, high) smoothly while behaving approximately like the identity mapping in the middle of the interval.

class homodyne.optimization.cmc.scaling.ParameterScaling

Bases: object

Scaling parameters for a single parameter.

name

Parameter name.

Type:

str

center

Center value for transformation (typically prior mean or bounds midpoint).

Type:

float

scale

Scale factor for transformation (typically prior std or bounds/4).

Type:

float

low

Lower bound for clipping.

Type:

float

high

Upper bound for clipping.

Type:

float

name: str

center: float

scale: float

low: float

high: float

to_normalized(value)

Transform from original to normalized space.

Uses the analytic inverse of the smooth bounding transform to recover the underlying affine value prior to normalization.

Return type:

float

to_original(z_value)

Transform from normalized to original space with smooth bounding.

Return type:

Array

__init__(name, center, scale, low, high)

homodyne.optimization.cmc.scaling.compute_scaling_factors(parameter_space, n_phi, analysis_mode)

Compute scaling factors for all parameters.

Parameters:

• parameter_space (ParameterSpace) – Parameter space with bounds and priors.

• n_phi (int) – Number of phi angles.

• analysis_mode (str) – Analysis mode (“static” or “laminar_flow”).

Returns:

Scaling factors for each parameter.

Return type:

dict[str, ParameterScaling]

homodyne.optimization.cmc.scaling.sample_scaled_parameter(name, scaling, initial_z=None, prior_scale=1.0)

Sample a parameter in normalized space and transform to original.

Parameters:

• name (str) – Parameter name (used for NumPyro site name).

• scaling (ParameterScaling) – Scaling parameters.

• initial_z (float | None) – Initial value in normalized space (for initialization).

• prior_scale (float) – Prior tempering scale factor. For CMC with K shards, set to sqrt(K) to implement prior^(1/K) tempering (Scott et al. 2016). The z-space prior Normal(0, 1) becomes Normal(0, prior_scale), effectively widening the prior so the combined posterior across K shards has the correct single-prior contribution.

Returns:

Parameter value in original space.

Return type:

Array

homodyne.optimization.cmc.scaling.log_scaling_factors(scalings)

Log scaling factors for debugging.

Parameters:

scalings (dict[str, ParameterScaling]) – Scaling factors.

Return type:

None

homodyne.optimization.cmc.scaling.transform_initial_values_to_z(initial_values, scalings)

Transform initial values from original to normalized space.

Parameters:

• initial_values (dict[str, float] | None) – Initial values in original space.

• scalings (dict[str, ParameterScaling]) – Scaling factors.

Returns:

Initial values in normalized (z) space.

Return type:

dict[str, float]

homodyne.optimization.cmc.scaling.transform_samples_from_z(samples, scalings)

Transform samples from normalized to original space.

Parameters:

• samples (dict[str, Array]) – Samples in normalized space (keys ending with “_z”).

• scalings (dict[str, ParameterScaling]) – Scaling factors.

Returns:

Samples in original space.

Return type:

dict[str, Array]

Understanding Z-Space Parameters

CMC uses non-centered parameterization to balance gradient magnitudes across parameters with vastly different scales (e.g., D0 ~ 10^4 vs gamma_dot_t0 ~ 10^-3).

When sampling, parameters are transformed to normalized z-space:

• Each parameter is sampled as z ~ Normal(0, 1)

• Transformed to original space: param = center + scale * z

• Clipped to physical bounds

MCMC Sample Names:

The MCMC output includes both z-space and original-space parameter names:

Z-Space Name	Original Name	Description
D0_z	D0	Diffusion coefficient (normalized / original)
alpha_z	alpha	Diffusion exponent
contrast_0_z	contrast_0	Per-angle contrast (phi index 0)
offset_0_z	offset_0	Per-angle offset (phi index 0)

Filtering Samples:

When working with MCMC samples, you may want to filter out z-space parameters:

# Get only original-space parameters
original_params = {k: v for k, v in samples.items() if not k.endswith('_z')}

# Get only physical parameters (exclude sigma, n_numerical_issues)
physical_params = ['D0', 'alpha', 'D_offset', 'gamma_dot_t0', 'beta',
                   'gamma_dot_t_offset', 'phi0']
physical_samples = {k: v for k, v in samples.items() if k in physical_params}

Results

CMC result dataclass and ArviZ integration.

This module provides the CMCResult dataclass that encapsulates MCMC posterior samples and diagnostics in a format compatible with ArviZ and the existing CLI save functions.

class homodyne.optimization.cmc.results.ParameterStats

Bases: dict

Hybrid mapping/sequence for posterior summaries.

Supports dict-style access by name (for tests/back-compat) and list/array-style access by index (for plotting utilities).

__init__(ordered_names, values)

property as_array: ndarray

Return ordered values as numpy array.

tolist()

Return ordered values as list (numpy compatibility).

Return type:

list[float]

class homodyne.optimization.cmc.results.CMCResult

Bases: object

CMC analysis result with posterior samples and diagnostics.

This dataclass is compatible with save_mcmc_results() in cli/commands.py.

parameters

Posterior mean values, shape (n_params,).

Type:

np.ndarray

uncertainties

Posterior standard deviations, shape (n_params,).

Type:

np.ndarray

param_names

Parameter names in sampling order.

Type:

list[str]

samples

Raw samples, {name: (n_chains, n_samples)}.

Type:

dict[str, np.ndarray]

convergence_status

“converged” | “divergences” | “not_converged”.

Type:

str

r_hat

Per-parameter R-hat values.

Type:

dict[str, float]

ess_bulk

Per-parameter bulk ESS.

Type:

dict[str, float]

ess_tail

Per-parameter tail ESS.

Type:

dict[str, float]

divergences

Total number of divergent transitions.

Type:

int

inference_data

ArviZ InferenceData for plotting.

Type:

az.InferenceData

execution_time

Total sampling time in seconds.

Type:

float

warmup_time

Warmup time in seconds.

Type:

float

n_chains

Number of MCMC chains.

Type:

int

n_samples

Samples per chain.

Type:

int

n_warmup

Warmup samples.

Type:

int

analysis_mode

Analysis mode used.

Type:

str

covariance

Parameter covariance matrix (from samples).

Type:

np.ndarray

chi_squared

Placeholder for compatibility (not directly computed in MCMC).

Type:

float

reduced_chi_squared

Placeholder for compatibility.

Type:

float

device_info

Device used for computation.

Type:

dict[str, Any]

parameters: ndarray

uncertainties: ndarray

param_names: list[str]

samples: dict[str, ndarray]

convergence_status: str

r_hat: dict[str, float]

ess_bulk: dict[str, float]

ess_tail: dict[str, float]

divergences: int

inference_data: arviz.InferenceData

execution_time: float

warmup_time: float

n_chains: int = 4

n_samples: int = 2000

n_warmup: int = 500

analysis_mode: str = 'static'

per_angle_mode: str = 'auto'

num_shards: int = 1

covariance: ndarray

chi_squared: float = 0.0

reduced_chi_squared: float = 0.0

device_info: dict[str, Any]

recovery_actions: list[str]

quality_flag: str = 'good'

mean_params: ParameterStats

std_params: ParameterStats

mean_contrast: float | None = None

std_contrast: float | None = None

mean_offset: float | None = None

std_offset: float | None = None

is_cmc_result()

Return True - required by CLI for diagnostic generation.

Return type:

bool

property success: bool

Return True if sampling converged (backward compatibility).

property message: str

Return descriptive message about result.

classmethod from_mcmc_samples(mcmc_samples, stats, analysis_mode, n_warmup=500, min_ess=None)

Create CMCResult from MCMC samples.

Parameters:

• mcmc_samples (MCMCSamples) – Raw MCMC samples.

• stats (SamplingStats) – Sampling statistics.

• analysis_mode (str) – Analysis mode used.

• n_warmup (int) – Number of warmup samples.

• min_ess (float | None) – Minimum effective sample size for convergence checks. If None, uses DEFAULT_MIN_ESS from diagnostics module.

Returns:

Complete result object.

Return type:

CMCResult

get_posterior_stats()

Get posterior statistics for each parameter.

Returns:

Statistics per parameter: mean, std, median, hdi_5%, hdi_95%.

Return type:

dict[str, dict[str, float]]

get_samples_array()

Get samples as 3D array.

Returns:

Shape (n_chains, n_samples, n_params).

Return type:

ndarray

validate_parameters(n_phi=None)

Validate that result contains expected parameters.

Parameters:

n_phi (int | None) – Number of phi angles expected. If None, infers from samples.

Returns:

List of validation warnings (empty if all valid).

Return type:

list[str]

__init__(parameters, uncertainties, param_names, samples, convergence_status, r_hat, ess_bulk, ess_tail, divergences, inference_data, execution_time, warmup_time, n_chains=4, n_samples=2000, n_warmup=500, analysis_mode='static', per_angle_mode='auto', num_shards=1, covariance=<factory>, chi_squared=0.0, reduced_chi_squared=0.0, device_info=<factory>, recovery_actions=<factory>, quality_flag='good', mean_params=<factory>, std_params=<factory>, mean_contrast=None, std_contrast=None, mean_offset=None, std_offset=None)

homodyne.optimization.cmc.results.create_inference_data(mcmc_samples)

Create ArviZ InferenceData from MCMC samples.

Parameters:

mcmc_samples (MCMCSamples) – Raw MCMC samples.

Returns:

ArviZ-compatible data structure.

Return type:

InferenceData

homodyne.optimization.cmc.results.samples_dict_from_array(samples_array, param_names)

Convert samples array to dictionary.

Parameters:

• samples_array (ndarray) – Shape (n_chains, n_samples, n_params).

• param_names (list[str]) – Parameter names.

Returns:

Samples dictionary.

Return type:

dict[str, ndarray]

homodyne.optimization.cmc.results.compute_fitted_c2(result, t1, t2, phi, q, L, dt, analysis_mode, fixed_contrasts=None, fixed_offsets=None)

Compute fitted C2 values from posterior mean.

Parameters:

• result (CMCResult) – CMC result with posterior samples.

• t1 (ndarray) – Coordinates (pooled 1D).

• t2 (ndarray) – Coordinates (pooled 1D).

• phi (ndarray) – Coordinates (pooled 1D).

• q (float) – Physics parameters.

• L (float) – Physics parameters.

• dt (float) – Physics parameters.

• analysis_mode (str) – Analysis mode.

• fixed_contrasts (ndarray | None) – Per-angle contrast array of shape (n_phi,) for constant and constant_averaged modes where contrast is not sampled. Required when neither contrast_0 nor contrast appears in posterior samples.

• fixed_offsets (ndarray | None) – Per-angle offset array of shape (n_phi,) paired with fixed_contrasts.

Returns:

(c2_fitted_mean, c2_fitted_std) from posterior.

Return type:

tuple[ndarray, ndarray]

Diagnostics

MCMC convergence diagnostics including R-hat, effective sample size (ESS), and divergence analysis.

Convergence diagnostics for CMC analysis.

This module provides functions for computing MCMC convergence diagnostics including R-hat, effective sample size (ESS), and divergence checks.

homodyne.optimization.cmc.diagnostics.compute_r_hat(samples)

Compute split-R-hat (Vehtari et al. 2021) for each parameter.

Uses ArviZ’s implementation of split-R-hat, which splits each chain in half before computing R-hat across 2*n_chains half-chains. This detects both between-chain discordance and within-chain non-stationarity that the original 1992 Gelman-Rubin formula misses.

Falls back to the classical Gelman-Rubin formula when ArviZ is not available.

Parameters:

samples (dict[str, ndarray]) – Parameter samples, {name: (n_chains, n_samples)}.

Returns:

R-hat value for each parameter.

Return type:

dict[str, float]

homodyne.optimization.cmc.diagnostics.compute_ess(samples)

Compute effective sample size (bulk and tail) for each parameter.

ESS measures the number of independent samples accounting for autocorrelation. Higher is better.

Parameters:

samples (dict[str, ndarray]) – Parameter samples, {name: (n_chains, n_samples)}.

Returns:

(ess_bulk, ess_tail) dictionaries.

Return type:

tuple[dict[str, float], dict[str, float]]

homodyne.optimization.cmc.diagnostics.check_convergence(r_hat, ess_bulk, divergences, n_samples, n_chains, max_rhat=DEFAULT_MAX_RHAT, min_ess=DEFAULT_MIN_ESS, max_divergence_rate=DEFAULT_MAX_DIVERGENCE_RATE, num_shards=1)

Check convergence and generate warnings.

Parameters:

• r_hat (dict[str, float]) – Per-parameter R-hat values.

• ess_bulk (dict[str, float]) – Per-parameter bulk ESS values.

• divergences (int) – Number of divergent transitions.

• n_samples (int) – Samples per chain.

• n_chains (int) – Number of chains.

• max_rhat (float) – Maximum acceptable R-hat.

• min_ess (float) – Minimum acceptable ESS.

• max_divergence_rate (float) – Maximum acceptable divergence rate.

• num_shards (int) – Number of shards (for CMC). Divergences are summed across shards, so total transitions = num_shards × n_chains × n_samples.

Returns:

(status, warnings) where status is “converged” | “divergences” | “not_converged”.

Return type:

tuple[str, list[str]]

homodyne.optimization.cmc.diagnostics.create_diagnostics_dict(r_hat, ess_bulk, ess_tail, divergences, convergence_status, warnings, n_chains, n_warmup, n_samples, warmup_time, sampling_time, num_shards=1)

Create diagnostics dictionary for JSON output.

Parameters:

• r_hat (dict[str, float]) – Per-parameter R-hat.

• ess_bulk (dict[str, float]) – Per-parameter bulk ESS.

• ess_tail (dict[str, float]) – Per-parameter tail ESS.

• divergences (int) – Number of divergences.

• convergence_status (str) – Convergence status.

• warnings (list[str]) – Warning messages.

• n_chains (int) – Number of chains.

• n_warmup (int) – Warmup samples.

• n_samples (int) – Posterior samples.

• warmup_time (float) – Warmup time in seconds.

• sampling_time (float) – Sampling time in seconds.

• num_shards (int) – Number of shards combined (default 1). For CMC runs, divergences is the aggregate total across all shards, so the correct denominator is num_shards * n_chains * n_samples.

Returns:

Diagnostics dictionary.

Return type:

dict[str, Any]

homodyne.optimization.cmc.diagnostics.summarize_diagnostics(r_hat, ess_bulk, divergences, n_samples, n_chains, num_shards=1)

Create human-readable diagnostics summary.

Parameters:

• r_hat (dict[str, float]) – R-hat values.

• ess_bulk (dict[str, float]) – ESS values.

• divergences (int) – Divergence count.

• n_samples (int) – Samples per chain.

• n_chains (int) – Number of chains.

• num_shards (int) – Number of shards (for CMC).

Returns:

Summary string.

Return type:

str

homodyne.optimization.cmc.diagnostics.log_analysis_summary(convergence_status, r_hat, ess_bulk, divergences, n_samples, n_chains, n_shards, shards_succeeded, execution_time)

Log a comprehensive summary at the end of CMC analysis.

Parameters:

• convergence_status (str) – Final convergence status.

• r_hat (dict[str, float]) – Per-parameter R-hat values.

• ess_bulk (dict[str, float]) – Per-parameter bulk ESS values.

• divergences (int) – Total divergent transitions.

• n_samples (int) – Samples per chain.

• n_chains (int) – Number of chains.

• n_shards (int) – Total number of shards.

• shards_succeeded (int) – Number of successful shards.

• execution_time (float) – Total execution time in seconds.

Return type:

None

homodyne.optimization.cmc.diagnostics.get_convergence_recommendations(max_rhat, min_ess, divergences, n_samples, n_chains, num_shards=1)

Generate specific recommendations for convergence issues.

Parameters:

• max_rhat (float) – Maximum R-hat value across parameters.

• min_ess (float) – Minimum bulk ESS across parameters.

• divergences (int) – Number of divergent transitions.

• n_samples (int) – Samples per chain.

• n_chains (int) – Number of chains.

• num_shards (int) – Number of shards (for CMC).

Returns:

List of recommendation strings.

Return type:

list[str]

homodyne.optimization.cmc.diagnostics.compute_posterior_contraction(posterior_std, prior_std)

Compute Posterior Contraction Ratio (PCR).

PCR measures how much the data informed the posterior relative to the prior. PCR = 1 - (posterior_std / prior_std)

Interpretation: - PCR ≈ 0: Posterior ≈ prior (data didn’t constrain the parameter) - PCR ≈ 0.5: Posterior half as wide as prior (moderate constraint) - PCR ≈ 0.9: Posterior 10% as wide as prior (strong constraint) - PCR < 0: Posterior wider than prior (model misspecification or numerical issues)

Parameters:

• posterior_std (float) – Standard deviation of the posterior distribution.

• prior_std (float) – Standard deviation of the prior distribution.

Returns:

Posterior contraction ratio, typically in [0, 1].

Return type:

float

homodyne.optimization.cmc.diagnostics.compute_nlsq_comparison_metrics(cmc_mean, cmc_std, nlsq_value, nlsq_std=None)

Compute metrics comparing CMC posterior to NLSQ point estimate.

Parameters:

• cmc_mean (float) – CMC posterior mean.

• cmc_std (float) – CMC posterior standard deviation.

• nlsq_value (float) – NLSQ point estimate.

• nlsq_std (float | None) – NLSQ standard error. If None, only CMC-based metrics computed.

Returns:

Dictionary with comparison metrics: - z_score: abs(CMC_mean - NLSQ) / CMC_std (should be < 2 for consistency) - uncertainty_ratio: CMC_std / NLSQ_std (should be < 5x ideally) - relative_diff: (CMC_mean - NLSQ) / abs(NLSQ) (percent difference) - coverage: Whether NLSQ falls within CMC 95% CI

Return type:

dict[str, float]

homodyne.optimization.cmc.diagnostics.compute_precision_analysis(cmc_result, nlsq_result=None, nlsq_uncertainties=None, prior_stds=None)

Compute comprehensive precision analysis for all parameters.

Parameters:

• cmc_result (dict[str, dict]) – CMC posterior statistics, keyed by parameter name. Each entry should have “mean” and “std” keys.

• nlsq_result (dict[str, float] | None) – NLSQ point estimates, keyed by parameter name.

• nlsq_uncertainties (dict[str, float] | None) – NLSQ standard errors, keyed by parameter name.

• prior_stds (dict[str, float] | None) – Prior standard deviations, keyed by parameter name.

Returns:

Precision metrics for each parameter.

Return type:

dict[str, dict[str, float]]

homodyne.optimization.cmc.diagnostics.log_precision_analysis(analysis, log_fn=None, tolerance_pct=20.0)

Log a comprehensive precision analysis report.

Parameters:

• analysis (dict[str, dict[str, float]]) – Output from compute_precision_analysis().

• log_fn (Callable[[str], None] | None) – Logging function. If None, uses module logger.

• tolerance_pct (float) – Percentage tolerance threshold for flagging parameters. Default 20% - parameters exceeding this are flagged.

Returns:

Formatted analysis report.

Return type:

str

class homodyne.optimization.cmc.diagnostics.BimodalResult

Bases: object

Result of bimodal detection for a single parameter.

is_bimodal

Whether the posterior appears bimodal.

Type:

bool

weights

Component weights from GMM.

Type:

tuple[float, float]

means

Component means from GMM.

Type:

tuple[float, float]

stds

Per-component standard deviations from GMM.

Type:

tuple[float, float]

separation

Absolute distance between means.

Type:

float

relative_separation

Separation relative to scale (separation / |mean(means)|).

Type:

float

is_bimodal: bool

weights: tuple[float, float]

means: tuple[float, float]

stds: tuple[float, float]

separation: float

relative_separation: float

__init__(is_bimodal, weights, means, stds, separation, relative_separation)

class homodyne.optimization.cmc.diagnostics.ModeCluster

Bases: object

A single mode from bimodal consensus combination.

mean

Per-parameter consensus mean for this mode.

Type:

dict[str, float]

std

Per-parameter consensus std for this mode.

Type:

dict[str, float]

weight

Fraction of shards supporting this mode (0-1).

Type:

float

n_shards

Number of shards in this cluster.

Type:

int

samples

Generated samples from N(mean, std^2), shape (n_chains, n_samples).

Type:

dict[str, np.ndarray]

mean: dict[str, float]

std: dict[str, float]

weight: float

n_shards: int

samples: dict[str, ndarray]

__init__(mean, std, weight, n_shards, samples)

class homodyne.optimization.cmc.diagnostics.BimodalConsensusResult

Bases: object

Result of mode-aware consensus combination.

Attached to MCMCSamples when bimodal posteriors are detected and per-mode consensus is used instead of standard combination.

modes

Mode clusters (typically 2) with per-mode consensus statistics.

Type:

list[ModeCluster]

modal_params

Parameter names that triggered bimodal detection.

Type:

list[str]

co_occurrence

Cross-parameter co-occurrence info (e.g., D0-alpha correlation).

Type:

dict[str, Any]

modes: list[ModeCluster]

modal_params: list[str]

co_occurrence: dict[str, Any]

__init__(modes, modal_params, co_occurrence)

homodyne.optimization.cmc.diagnostics.detect_bimodal(samples, min_weight=0.2, min_relative_separation=0.5)

Detect bimodality using 2-component Gaussian Mixture Model.

Parameters:

• samples (ndarray) – 1D array of posterior samples.

• min_weight (float) – Minimum weight for both components to be considered bimodal.

• min_relative_separation (float) – Minimum separation between means (relative to scale) for bimodality.

Returns:

Detection result with component details.

Return type:

BimodalResult

homodyne.optimization.cmc.diagnostics.check_shard_bimodality(samples, params_to_check=None)

Check multiple parameters for bimodality.

Parameters:

• samples (dict[str, ndarray]) – Parameter samples from a shard.

• params_to_check (list[str] | None) – Parameters to check. Defaults to key physical parameters.

Returns:

Mapping from param name to BimodalResult.

Return type:

dict[str, BimodalResult]

homodyne.optimization.cmc.diagnostics.summarize_cross_shard_bimodality(bimodal_detections, n_shards, consensus_means=None, significance_threshold=0.05)

Aggregate per-shard bimodal detections into a cross-shard summary.

Groups detections by parameter, computes mode statistics, separation significance, and D0-alpha co-occurrence to quantify consensus distortion.

Parameters:

• bimodal_detections (list[dict[str, Any]]) – Per-detection records, each with keys: “shard”, “param”, “mode1”, “mode2”, “weights”, “separation”.

• n_shards (int) – Total number of successful shards (denominator for bimodal fraction).

• consensus_means (dict[str, float] | None) – Mean-of-means for each parameter (pre-combine estimate). Used to check if consensus falls in a density trough between modes.

• significance_threshold (float) – Minimum bimodal fraction (detections/n_shards) to include a parameter in the summary. Default 5%.

Returns:

Summary with keys: - “per_param”: dict mapping param name to per-parameter stats - “co_occurrence”: dict with D0-alpha co-occurrence info - “n_detections”: total detection count - “n_shards”: total shard count

Return type:

dict[str, Any]

homodyne.optimization.cmc.diagnostics.cluster_shard_modes(bimodal_detections, successful_samples, bimodal_summary, param_bounds)

Jointly cluster shards into two mode populations.

Uses range-normalized feature vectors from modal parameters to assign each shard to the nearest mode centroid. Bimodal shards contribute one component to each cluster.

Parameters:

• bimodal_detections (list[dict[str, Any]]) – Per-detection records with keys: “shard”, “param”, “mode1”, “mode2”, “std1”, “std2”, “weights”, “separation”.

• successful_samples (list[Any]) – List of MCMCSamples (or similar with .samples dict attribute).

• bimodal_summary (dict[str, Any]) – Output from summarize_cross_shard_bimodality().

• param_bounds (dict[str, tuple[float, float]]) – Parameter bounds for range-based normalization, {param: (lo, hi)}.

Returns:

(cluster_0_shards, cluster_1_shards) where cluster_0 is “lower” and cluster_1 is “upper”. Bimodal shards appear in both lists.

Return type:

tuple[list[int], list[int]]

Key Functions

homodyne.optimization.cmc.diagnostics.compute_r_hat	Compute split-R-hat (Vehtari et al. 2021) for each parameter.
homodyne.optimization.cmc.diagnostics.compute_ess	Compute effective sample size (bulk and tail) for each parameter.
homodyne.optimization.cmc.diagnostics.check_convergence	Check convergence and generate warnings.
homodyne.optimization.cmc.diagnostics.create_diagnostics_dict	Create diagnostics dictionary for JSON output.
homodyne.optimization.cmc.diagnostics.summarize_diagnostics	Create human-readable diagnostics summary.
homodyne.optimization.cmc.diagnostics.log_analysis_summary	Log a comprehensive summary at the end of CMC analysis.
homodyne.optimization.cmc.diagnostics.get_convergence_recommendations	Generate specific recommendations for convergence issues.

Convergence Thresholds

Default thresholds:

• MAX_RHAT: 1.05 (chains should have R-hat < 1.05 for convergence)

• MIN_ESS: 400 (effective sample size should exceed 400)

• MAX_DIVERGENCE_RATE: 5% (divergence rate should be < 5%)

Diagnostics Output:

The check_convergence function returns one of three statuses:

• converged: All chains mixed well, ESS adequate, no excessive divergences

• divergences: High divergence rate indicates model geometry issues

• not_converged: R-hat or ESS thresholds not met

I/O Operations

I/O utilities for CMC results.

This module provides functions for saving CMC results to files:

• samples.npz: ArviZ-compatible posterior samples

• fitted_data.npz: Fitted data matching NLSQ format

• parameters.json: Posterior statistics

• diagnostics.json: Convergence diagnostics

homodyne.optimization.cmc.io.save_samples_npz(result, output_path)

Save posterior samples in ArviZ-compatible format.

The saved file can be loaded directly with numpy and converted to ArviZ InferenceData without modification.

Parameters:

• result (CMCResult) – CMC result with samples.

• output_path (Path) – Output file path.

• Format (File)

• -----------

• schema_version (-)

• posterior_samples (-)

• param_names (-)

• r_hat (-)

• ess_bulk (-)

• ess_tail (-)

• divergences (-)

• analysis_mode (-)

• n_phi (-)

• n_chains (-)

• n_samples (-)

Return type:

None

homodyne.optimization.cmc.io.load_samples_npz(input_path)

Load samples from npz file.

Parameters:

input_path (Path) – Path to samples.npz file.

Returns:

Loaded data dictionary.

Return type:

dict[str, Any]

Raises:

• ValueError – If path validation fails (path traversal, non-existent file).

• FileNotFoundError – If the file does not exist.

homodyne.optimization.cmc.io.samples_to_arviz(samples_data)

Convert loaded samples to ArviZ InferenceData.

Parameters:

samples_data (dict[str, Any]) – Data from load_samples_npz().

Returns:

ArviZ-compatible data structure.

Return type:

az.InferenceData

homodyne.optimization.cmc.io.save_fitted_data_npz(result, c2_exp, c2_fitted, c2_fitted_std, t1, t2, phi_angles, q, output_path)

Save fitted data in NLSQ-compatible format.

Parameters:

• result (CMCResult) – CMC result.

• c2_exp (-) – Experimental C2 data.

• c2_fitted (-) – Fitted C2 (posterior mean).

• c2_fitted_std (-) – Fitted C2 uncertainty.

• t1 (-) – Time coordinates t1.

• t2 (Time coordinates) – Time coordinates t2.

• phi_angles (-) – Phi angles.

• q (-) – Wavevector.

• output_path (Path) – Output file path.

• Format (File)

• -----------

• c2_exp

• c2_fitted

• residuals (-)

• c2_fitted_std

• c2_fitted_5pct (-)

• c2_fitted_95pct (-)

• q

• phi_angles

• t1

• t2

Return type:

None

homodyne.optimization.cmc.io.save_parameters_json(result, output_path)

Save posterior parameter statistics to JSON.

Parameters:

• result (CMCResult) – CMC result.

• output_path (Path) – Output file path.

Return type:

None

homodyne.optimization.cmc.io.save_diagnostics_json(result, output_path, warnings=None)

Save convergence diagnostics to JSON.

Parameters:

• result (CMCResult) – CMC result.

• output_path (Path) – Output file path.

• warnings (list[str] | None) – Warning messages from convergence check.

Return type:

None

homodyne.optimization.cmc.io.save_all_results(result, output_dir, c2_exp=None, c2_fitted=None, c2_fitted_std=None, t1=None, t2=None, phi_angles=None, q=None)

Save all CMC result files.

Parameters:

• result (CMCResult) – CMC result.

• output_dir (Path) – Output directory.

• c2_exp (ndarray | None) – Data for fitted_data.npz.

• c2_fitted (ndarray | None) – Data for fitted_data.npz.

• c2_fitted_std (ndarray | None) – Data for fitted_data.npz.

• t1 (ndarray | None) – Coordinates.

• t2 (ndarray | None) – Coordinates.

• phi_angles (ndarray | None) – Coordinates.

• q (float | None) – Wavevector.

Returns:

Paths to saved files.

Return type:

dict[str, Path]

Plotting

ArviZ diagnostic plots for CMC results.

This module provides the 6 standard ArviZ diagnostic plots: 1. Pair plot (corner plot) 2. Forest plot 3. Energy plot 4. Autocorrelation plot 5. Rank plot 6. ESS plot

homodyne.optimization.cmc.plotting.generate_diagnostic_plots(result, output_dir, figsize=DEFAULT_FIGSIZE, dpi=DEFAULT_DPI, param_subset=None)

Generate all 6 ArviZ diagnostic plots.

Parameters:

• result (CMCResult) – CMC result with inference_data.

• output_dir (Path) – Directory to save plots.

• figsize (tuple[int, int]) – Figure size in inches.

• dpi (int) – Figure resolution.

• param_subset (list[str] | None) – Subset of parameters to plot. If None, plots all.

Returns:

Paths to saved plot files.

Return type:

list[Path]

homodyne.optimization.cmc.plotting.plot_pair(idata, output_dir, var_names=None, figsize=DEFAULT_FIGSIZE, dpi=DEFAULT_DPI)

Generate pair (corner) plot.

Shows pairwise parameter correlations and marginal distributions.

Parameters:

• idata (InferenceData) – ArviZ inference data.

• output_dir (Path) – Output directory.

• var_names (list[str] | None) – Parameters to include.

• figsize (tuple[int, int]) – Figure size.

• dpi (int) – Resolution.

Returns:

Path to saved plot.

Return type:

Path

homodyne.optimization.cmc.plotting.plot_forest(idata, output_dir, var_names=None, figsize=DEFAULT_FIGSIZE, dpi=DEFAULT_DPI)

Generate forest plot.

Shows posterior distributions with HDI intervals.

Parameters:

• idata (InferenceData) – ArviZ inference data.

• output_dir (Path) – Output directory.

• var_names (list[str] | None) – Parameters to include.

• figsize (tuple[int, int]) – Figure size.

• dpi (int) – Resolution.

Returns:

Path to saved plot.

Return type:

Path

homodyne.optimization.cmc.plotting.plot_energy(idata, output_dir, figsize=(10, 6), dpi=DEFAULT_DPI)

Generate energy plot.

Compares marginal energy distribution to energy transition distribution. Large differences indicate sampling problems.

Parameters:

• idata (InferenceData) – ArviZ inference data.

• output_dir (Path) – Output directory.

• figsize (tuple[int, int]) – Figure size.

• dpi (int) – Resolution.

Returns:

Path to saved plot.

Return type:

Path

homodyne.optimization.cmc.plotting.plot_autocorr(idata, output_dir, var_names=None, figsize=DEFAULT_FIGSIZE, dpi=DEFAULT_DPI)

Generate autocorrelation plot.

Shows how quickly samples become independent.

Parameters:

• idata (InferenceData) – ArviZ inference data.

• output_dir (Path) – Output directory.

• var_names (list[str] | None) – Parameters to include.

• figsize (tuple[int, int]) – Figure size.

• dpi (int) – Resolution.

Returns:

Path to saved plot.

Return type:

Path

homodyne.optimization.cmc.plotting.plot_rank(idata, output_dir, var_names=None, figsize=DEFAULT_FIGSIZE, dpi=DEFAULT_DPI)

Generate rank plot.

Rank plots help identify chain mixing problems.

Parameters:

• idata (InferenceData) – ArviZ inference data.

• output_dir (Path) – Output directory.

• var_names (list[str] | None) – Parameters to include.

• figsize (tuple[int, int]) – Figure size.

• dpi (int) – Resolution.

Returns:

Path to saved plot.

Return type:

Path

homodyne.optimization.cmc.plotting.plot_ess(idata, output_dir, var_names=None, figsize=(10, 6), dpi=DEFAULT_DPI)

Generate ESS evolution plot.

Shows how effective sample size grows with more samples.

Parameters:

• idata (InferenceData) – ArviZ inference data.

• output_dir (Path) – Output directory.

• var_names (list[str] | None) – Parameters to include.

• figsize (tuple[int, int]) – Figure size.

• dpi (int) – Resolution.

Returns:

Path to saved plot.

Return type:

Path

homodyne.optimization.cmc.plotting.plot_trace(idata, output_dir, var_names=None, figsize=DEFAULT_FIGSIZE, dpi=DEFAULT_DPI)

Generate trace plot (bonus diagnostic).

Shows parameter values over sampling iterations.

Parameters:

• idata (InferenceData) – ArviZ inference data.

• output_dir (Path) – Output directory.

• var_names (list[str] | None) – Parameters to include.

• figsize (tuple[int, int]) – Figure size.

• dpi (int) – Resolution.

Returns:

Path to saved plot.

Return type:

Path

Backends

CMC supports multiple parallelization backends for distributed MCMC execution.

CMC execution backends.

This module provides different execution backends for running CMC shards in parallel:

• MultiprocessingBackend: CPU parallelism via Python multiprocessing

• PjitBackend: JAX distributed execution via pjit

• PBSBackend: HPC cluster execution via PBS job scheduler

Backend Selection

Available backends:

• multiprocessing: Python multiprocessing for multi-core workstations (default)

• pjit: JAX pjit for single-node multi-device parallelism

• pbs: PBS job scheduler for HPC clusters

Backend Configuration:

The backend is auto-selected based on environment, but can be overridden via configuration:

optimization:
  cmc:
    sharding:
      backend: multiprocessing  # or pjit, pbs

Per-Shard Sampling Behavior

All backends follow the same per-shard sampling pattern:

1. Shard preparation: Extract data subset with associated metadata (phi indices, time arrays)

2. Model kwargs construction: Build model arguments for the shard’s data

3. Sampler invocation: Call run_nuts_sampling() with shard-specific data

4. Result collection: Gather MCMCSamples and SamplingStats

Per-shard execution (simplified from backends/multiprocessing.py):

def _run_single_shard(shard_data, config, model, ...):
    # Build model kwargs for this shard
    model_kwargs = {
        "data": shard_data.data,
        "t1": shard_data.t1,
        "t2": shard_data.t2,
        "phi_indices": shard_data.phi_indices,
        ...
    }

    # Same sampler as single-shard path
    samples, stats = run_nuts_sampling(
        model=model,
        model_kwargs=model_kwargs,
        config=config,
        initial_values=initial_values,
        ...
    )
    return samples, stats

What each shard receives:

• Subset of data points (respecting max_points_per_shard)

• Full phi_unique array (all angles, for proper indexing)

• Shard-specific phi_indices (mapping points to angles)

• Same physics parameters (q, L, dt, time_grid)

• Same MCMC configuration (num_warmup, num_samples, etc.)

What each shard produces:

• MCMCSamples: Posterior samples for all parameters

• SamplingStats: Timing, divergences, acceptance rate

• Per-shard diagnostics: R-hat, ESS (within-shard convergence)

The combination phase (see Sharding Strategy (Detailed)) then merges these independent subposteriors using precision-weighted Gaussian consensus.

Base Backend

Base class for CMC execution backends.

This module defines the abstract interface for CMC backends and provides a factory function for selecting backends.

class homodyne.optimization.cmc.backends.base.CMCBackend

Bases: ABC

Abstract base class for CMC execution backends.

Backends handle the parallel execution of MCMC sampling across data shards and the combination of results.

abstractmethod run(model, model_kwargs, config, shards=None)

Run MCMC sampling (potentially across shards).

Parameters:

• model (Callable) – NumPyro model function.

• model_kwargs (dict[str, Any]) – Common model arguments.

• config (CMCConfig) – CMC configuration.

• shards (list[PreparedData] | None) – Data shards for parallel execution. If None, runs single-threaded on full data.

Returns:

Combined samples from all shards.

Return type:

MCMCSamples

abstractmethod get_name()

Get backend name.

Return type:

str

is_available()

Check if backend is available.

Returns:

True if backend can be used.

Return type:

bool

homodyne.optimization.cmc.backends.base.select_backend(config)

Select appropriate backend based on configuration.

Parameters:

config (CMCConfig) – CMC configuration.

Returns:

Selected backend instance.

Return type:

CMCBackend

Raises:

ValueError – If requested backend is not available.

homodyne.optimization.cmc.backends.base.combine_shard_samples(shard_samples, method='weighted_gaussian', chunk_size=500)

Combine samples from multiple shards.

For K <= chunk_size shards, uses a single-pass combination.

For K > chunk_size shards (hierarchical mode), accumulates posterior moments (mean, variance) across chunks without drawing intermediate synthetic samples. A single Gaussian draw is performed at the end from the aggregated moments. This avoids the precision-multiplication artefact that arises when recursive combination re-applies precision-weighting to synthetically drawn intermediate samples (P1-R6-01).

Memory scaling:

• Each shard result: ~100KB (13 params x 4 chains x 1500 samples x 8 bytes)

• Hierarchical (chunk=500): processes max(chunk_size) shards at once (~50MB), then releases them. Moment accumulation uses O(n_params) space.

Parameters:

• shard_samples (list[MCMCSamples]) – Samples from each shard.

• method (str) – Combination method: “robust_consensus_mc” (recommended), “consensus_mc”, “weighted_gaussian”, “simple_average”, or “auto”.

• chunk_size (int) – Number of shards to process per chunk for hierarchical combination. Default 500 keeps peak memory under ~50MB per processing step.

Returns:

Combined samples.

Return type:

MCMCSamples

homodyne.optimization.cmc.backends.base.combine_shard_samples_bimodal(shard_samples, cluster_assignments, bimodal_detections, modal_params, co_occurrence, method='consensus_mc', chunk_seed=0)

Combine shard samples using mode-aware consensus.

For bimodal shards, uses per-component GMM statistics instead of full-posterior statistics to avoid density-trough corruption.

Parameters:

• shard_samples (list[MCMCSamples]) – All successful shard samples.

• cluster_assignments (tuple[list[int], list[int]]) – (lower_cluster_shards, upper_cluster_shards) from cluster_shard_modes(). Bimodal shards may appear in both lists.

• bimodal_detections (list[dict[str, Any]]) – Per-detection records with “shard”, “param”, “mode1”, “mode2”, “std1”, “std2”, “weights”.

• modal_params (list[str]) – Parameters that triggered bimodal detection.

• co_occurrence (dict[str, Any]) – Cross-parameter co-occurrence info.

• method (str) – Base combination method for non-modal params.

Returns:

(combined_samples, bimodal_result) where combined_samples has mixture-drawn primary samples and bimodal_result has per-mode details.

Return type:

tuple[MCMCSamples, BimodalConsensusResult]

Multiprocessing Backend

Multiprocessing backend for CMC execution.

This module provides parallel MCMC execution using Python’s multiprocessing module for CPU-based parallelism.

Optimizations (v2.9.1): - Batch PRNG key generation: Pre-generate all shard keys in single JAX call - Adaptive polling: Adjust poll interval based on shard activity - Event.wait heartbeat: Efficient heartbeat using Event.wait(timeout)

Optimizations (v2.22.2): - LPT scheduling: Dispatch highest-cost shards first (size + noise weighted) - Per-shard shared memory: Shard arrays stored in shared memory (avoids pickle overhead) - deque for pending shards: O(1) popleft instead of O(n) list.pop(0) - JIT cache fix: Enable persistent compilation cache via jax.config.update (env var alone insufficient in JAX 0.8+, min_compile_time lowered to 0)

class homodyne.optimization.cmc.backends.multiprocessing.SharedDataManager

Bases: object

Manages shared memory blocks for data common to all CMC shards.

Uses multiprocessing.shared_memory to share config, parameter space, initial values, and time_grid across spawned worker processes, avoiding redundant pickling per shard.

Note on serialization: Uses pickle internally for trusted config dicts only (CMCConfig.to_dict(), ParameterSpace). This matches the existing multiprocessing behavior which also pickles all process arguments.

Must be used as a context manager or call cleanup() in a finally block.

__init__()

create_shared_bytes(name, data)

Store bytes in shared memory.

Return type:

dict[str, Any]

create_shared_array(name, array)

Store a numpy array in shared memory.

Return type:

dict[str, Any]

create_shared_dict(name, d)

Serialize a trusted internal dict to shared memory.

Only used for CMCConfig and ParameterSpace dicts — never for external/untrusted data.

Return type:

dict[str, Any]

create_shared_shard_arrays(shard_data_list)

Place per-shard numpy arrays into shared memory (packed format).

Instead of creating one SharedMemory segment per array per shard (n_shards * 5 = thousands of file descriptors), this concatenates all shard arrays for each key into a single shared memory block. Only 5 SharedMemory segments are created regardless of shard count.

Parameters:

shard_data_list (list[dict[str, Any]]) – List of shard data dicts, each containing numpy arrays (data, t1, t2, phi_unique, phi_indices) and a scalar noise_scale.

Returns:

List of lightweight shard references (shm names + offsets). Each ref dict is small enough to serialize cheaply through spawn.

Return type:

list[dict[str, Any]]

cleanup()

Release all shared memory blocks. Must be called in a finally block.

Return type:

None

class homodyne.optimization.cmc.backends.multiprocessing.MultiprocessingBackend

Bases: CMCBackend

CMC backend using Python multiprocessing.

Runs MCMC sampling in parallel across CPU cores using Python’s multiprocessing module.

__init__(n_workers=None, spawn_method='spawn')

Initialize multiprocessing backend.

Parameters:

• n_workers (int | None) – Number of worker processes. If None, uses CPU count.

• spawn_method (str) – Process start method: “spawn”, “fork”, or “forkserver”.

get_name()

Get backend name.

Return type:

str

run(model, model_kwargs, config, shards=None, initial_values=None, parameter_space=None, analysis_mode='static', progress_bar=True)

Run MCMC sampling across shards.

Parameters:

• model (Callable) – NumPyro model function.

• model_kwargs (dict[str, Any]) – Common model arguments.

• config (CMCConfig) – CMC configuration.

• shards (list[PreparedData] | None) – Data shards.

• initial_values (dict[str, float] | None) – Initial parameter values.

• parameter_space (ParameterSpace | None) – Parameter space for priors.

• analysis_mode (str) – Analysis mode.

• progress_bar (bool) – Whether to show progress bar for shard completion.

Returns:

Combined samples from all shards.

Return type:

MCMCSamples

is_available()

Check if multiprocessing is available.

Return type:

bool

Key features:

• Automatic worker allocation based on CPU cores

• Configurable timeout handling

• Progress tracking with shard completion estimates

• Memory-efficient worker pool management

PJIT Backend

JAX pjit backend for CMC distributed execution.

This module provides distributed MCMC execution using JAX’s pjit for sharded computation across CPU devices.

Note: This is a CPU-only implementation per v2.3.0 architecture decision.

class homodyne.optimization.cmc.backends.pjit.PjitBackend

Bases: CMCBackend

JAX pjit backend for distributed MCMC execution.

Uses JAX’s pjit for parallel execution across CPU devices. This backend is suitable for multi-core CPU systems where JAX can leverage multiple devices.

Note

CPU-only per homodyne v2.3.0 architecture decision. For GPU support, use homodyne v2.2.1 or earlier.

__init__()

Initialize pjit backend.

get_name()

Get backend name.

Returns:

Backend identifier.

Return type:

str

is_available()

Check if pjit backend is available.

Returns:

True if JAX pjit can be used.

Return type:

bool

run(model, model_kwargs, config, shards=None, *, initial_values=None, parameter_space=None, analysis_mode=None, progress_bar=True)

Run MCMC sampling using pjit for parallelism.

Parameters:

• model (Callable) – NumPyro model function.

• model_kwargs (dict[str, Any]) – Common model arguments (q, L, dt, etc.).

• config (CMCConfig) – CMC configuration.

• shards (list[PreparedData] | None) – Data shards for parallel execution. If None, runs on full data without sharding.

Notes

Additional keyword arguments are accepted for signature compatibility with other backends (multiprocessing). They are currently unused but harmless, ensuring legacy calls with initial_values/parameter_space do not fail.

Returns:

Combined samples from all shards.

Return type:

MCMCSamples

PBS Backend

PBS (Portable Batch System) backend for CMC HPC cluster execution.

This module provides distributed MCMC execution on HPC clusters using PBS job scheduling.

Note: This backend requires: - PBS/Torque job scheduler (qsub, qstat commands) - Shared filesystem accessible from all nodes - homodyne installed on compute nodes

class homodyne.optimization.cmc.backends.pbs.PBSBackend

Bases: CMCBackend

PBS backend for HPC cluster MCMC execution.

Submits each data shard as a separate PBS job and combines results after all jobs complete.

Parameters:

• queue (str) – PBS queue name (default: “batch”).

• ppn (int) – Processors per node (default: 4).

• walltime (str) – Job walltime (default: “04:00:00”).

• memory (str) – Memory per job (default: “8gb”).

• poll_interval (int) – Seconds between job status checks (default: 30).

• max_wait_time (int) – Maximum wait time in seconds (default: 14400 = 4 hours).

__init__(queue='batch', ppn=4, walltime='04:00:00', memory='8gb', poll_interval=30, max_wait_time=14400)

Initialize PBS backend.

get_name()

Get backend name.

Returns:

Backend identifier.

Return type:

str

is_available()

Check if PBS backend is available.

Returns:

True if PBS commands are accessible.

Return type:

bool

Notes

P2-R6-06: Previously ran bare qsub with no arguments, which exits non-zero on all PBS/Torque versions (missing jobscript), so this method always returned False on valid clusters. Now checks for the presence of the qsub binary via shutil.which, which is sufficient to determine availability without triggering an error submission.

run(model, model_kwargs, config, shards=None)

Run MCMC sampling via PBS job submission.

Parameters:

• model (Callable) – NumPyro model function (not directly used - workers import it).

• model_kwargs (dict[str, Any]) – Common model arguments.

• config (CMCConfig) – CMC configuration.

• shards (list[PreparedData] | None) – Data shards for parallel execution.

Returns:

Combined samples from all PBS jobs.

Return type:

MCMCSamples

Raises:

RuntimeError – If jobs fail or timeout.

Anti-Degeneracy Defense System

The NLSQ module includes a comprehensive anti-degeneracy defense system for laminar flow analysis with many phi angles. See Anti-Degeneracy Defense System for theoretical background and usage tutorials.

Fourier Reparameterization (Layer 1)

Reduces per-angle parameter count by expressing contrast/offset as Fourier series.

Fourier Reparameterization for Anti-Degeneracy Defense.

This module replaces n_phi independent per-angle contrast/offset values with truncated Fourier series, dramatically reducing structural degeneracy.

Part of Anti-Degeneracy Defense System v2.9.0. See: docs/specs/anti-degeneracy-defense-v2.9.0.md

Mathematical Formulation

contrast(φ) = c₀ + Σₖ[cₖ×cos(kφ) + sₖ×sin(kφ)] for k=1..order offset(φ) = o₀ + Σₖ[oₖ×cos(kφ) + tₖ×sin(kφ)] for k=1..order

For order=2: - Contrast: 5 coefficients [c₀, c₁, s₁, c₂, s₂] - Offset: 5 coefficients [o₀, o₁, t₁, o₂, t₂] - Total: 10 Fourier coefficients vs 2×n_phi independent params

Parameter Count Comparison:

n_phi | Independent | Fourier (order=2) | Reduction
------|-------------|-------------------|----------
  2   |     4       |        4          |    0%
  3   |     6       |        6          |    0%
 10   |    20       |       10          |   50%
 23   |    46       |       10          |   78%
100   |   200       |       10          |   95%

Note: For n_phi <= 2*(order+1), independent mode is used.

class homodyne.optimization.nlsq.fourier_reparam.FourierReparamConfig

Bases: object

Configuration for Fourier reparameterization.

mode

Per-angle parameter mode: - “independent”: Use n_phi independent contrast/offset values - “fourier”: Use truncated Fourier series - “auto”: Use Fourier when n_phi > auto_threshold

Type:

str

fourier_order

Number of Fourier harmonics. Default 2. order=2 gives 5 coefficients per parameter (c0, c1, s1, c2, s2).

Type:

int

auto_threshold

Use Fourier when n_phi > this threshold in auto mode. Default 6.

Type:

int

c0_bounds

Bounds for mean contrast coefficient. Default (0.1, 0.8).

Type:

tuple

ck_bounds

Bounds for harmonic contrast amplitudes. Default (-0.2, 0.2).

Type:

tuple

o0_bounds

Bounds for mean offset coefficient. Default (0.5, 1.5).

Type:

tuple

ok_bounds

Bounds for harmonic offset amplitudes. Default (-0.3, 0.3).

Type:

tuple

mode: Literal['independent', 'fourier', 'auto'] = 'auto'

fourier_order: int = 2

auto_threshold: int = 6

c0_bounds: tuple[float, float] = (0.1, 0.8)

ck_bounds: tuple[float, float] = (-0.2, 0.2)

o0_bounds: tuple[float, float] = (0.5, 1.5)

ok_bounds: tuple[float, float] = (-0.3, 0.3)

classmethod from_dict(config_dict)

Create config from dictionary.

Return type:

FourierReparamConfig

__init__(mode='auto', fourier_order=2, auto_threshold=6, c0_bounds=(0.1, 0.8), ck_bounds=(-0.2, 0.2), o0_bounds=(0.5, 1.5), ok_bounds=(-0.3, 0.3))

class homodyne.optimization.nlsq.fourier_reparam.FourierReparameterizer

Bases: object

Handles conversion between Fourier coefficients and per-angle values.

This class provides the core functionality for Fourier reparameterization: 1. Convert per-angle values to Fourier coefficients (initialization) 2. Convert Fourier coefficients to per-angle values (model evaluation) 3. Compute Jacobian for covariance transformation

The Fourier basis ensures smooth variation of contrast/offset with angle, preventing the optimizer from using per-angle parameters to absorb angle-dependent physical signals (like the shear term cos(φ₀-φ)).

Parameters:

• phi_angles (ndarray) – Unique phi angles in radians, shape (n_phi,).

• config (FourierReparamConfig) – Fourier configuration.

n_phi

Number of unique phi angles.

Type:

int

n_coeffs

Total number of Fourier coefficients (contrast + offset).

Type:

int

n_coeffs_per_param

Coefficients per parameter type (contrast or offset).

Type:

int

use_fourier

Whether Fourier mode is active.

Type:

bool

Examples

>>> phi_angles = np.linspace(-np.pi, np.pi, 23)
>>> config = FourierReparamConfig(mode="fourier", fourier_order=2)
>>> fourier = FourierReparameterizer(phi_angles, config)
>>> # Convert initial per-angle values to Fourier
>>> contrast = np.full(23, 0.3)
>>> offset = np.full(23, 1.0)
>>> fourier_coeffs = fourier.per_angle_to_fourier(contrast, offset)
>>> # Convert back during model evaluation
>>> contrast_out, offset_out = fourier.fourier_to_per_angle(fourier_coeffs)

__init__(phi_angles, config)

Initialize Fourier reparameterizer.

Parameters:

• phi_angles (ndarray) – Unique phi angles in radians, shape (n_phi,).

• config (FourierReparamConfig) – Fourier configuration.

get_basis_matrix()

Get the Fourier basis matrix for covariance transformation.

Returns:

Basis matrix of shape (n_phi, n_coeffs_per_param) if in Fourier mode, None if in independent mode. The basis matrix B satisfies: per_angle_values = B @ fourier_coeffs

Return type:

ndarray | None

Notes

Used for transforming covariance from Fourier space to per-angle space: pcov_per_angle = B @ pcov_fourier @ B.T

property order: int

Get the Fourier order (number of harmonics).

Returns:

Fourier order from config.

Return type:

int

fourier_to_per_angle(fourier_coeffs)

Convert Fourier coefficients to per-angle contrast/offset.

Parameters:

fourier_coeffs (ndarray) – Shape (n_coeffs,) = [c0,c1,s1,c2,s2,…,o0,o1,t1,o2,t2,…].

Return type:

tuple[ndarray, ndarray]

Returns:

• contrast (np.ndarray) – Per-angle contrast values, shape (n_phi,).

• offset (np.ndarray) – Per-angle offset values, shape (n_phi,).

Raises:

ValueError – If fourier_coeffs has wrong shape.

per_angle_to_fourier(contrast, offset)

Convert per-angle values to Fourier coefficients.

Uses least squares fitting when n_phi > n_coeffs_per_param.

Parameters:

• contrast (ndarray) – Per-angle contrast values, shape (n_phi,).

• offset (ndarray) – Per-angle offset values, shape (n_phi,).

Returns:

Fourier coefficients, shape (n_coeffs,).

Return type:

ndarray

Raises:

ValueError – If contrast or offset has wrong shape.

get_jacobian_transform()

Get Jacobian of transformation: d(per_angle)/d(fourier).

Used for covariance transformation back to per-angle space:

Cov_per_angle = J @ Cov_fourier @ J.T

Returns:

Jacobian matrix of shape (2*n_phi, n_coeffs).

Return type:

ndarray

get_bounds()

Get bounds for Fourier coefficients.

Return type:

tuple[ndarray, ndarray]

Returns:

• lower (np.ndarray) – Lower bounds, shape (n_coeffs,).

• upper (np.ndarray) – Upper bounds, shape (n_coeffs,).

get_initial_coefficients(contrast_init, offset_init)

Get initial Fourier coefficients from initial values.

Parameters:

• contrast_init (float | ndarray) – Initial contrast (scalar for uniform, array for per-angle).

• offset_init (float | ndarray) – Initial offset (scalar for uniform, array for per-angle).

Returns:

Initial Fourier coefficients.

Return type:

ndarray

get_coefficient_labels()

Get parameter labels for Fourier coefficients.

Returns:

Parameter labels.

Return type:

list[str]

to_fourier(per_angle_values)

Convert a single per-angle array to Fourier coefficients.

Convenience method for transforming one group (contrast or offset) at a time, rather than both together.

Parameters:

per_angle_values (ndarray) – Per-angle values, shape (n_phi,).

Returns:

Fourier coefficients, shape (n_coeffs_per_param,).

Return type:

ndarray

Raises:

ValueError – If per_angle_values has wrong shape.

from_fourier(fourier_coeffs)

Convert Fourier coefficients to per-angle values for a single group.

Convenience method for transforming one group (contrast or offset) at a time, rather than both together.

Parameters:

fourier_coeffs (ndarray) – Fourier coefficients, shape (n_coeffs_per_param,).

Returns:

Per-angle values, shape (n_phi,).

Return type:

ndarray

Raises:

ValueError – If fourier_coeffs has wrong shape.

get_diagnostics()

Get Fourier reparameterization diagnostics.

Returns:

Diagnostic information.

Return type:

dict

homodyne.optimization.nlsq.fourier_reparam.create_fourier_model_wrapper(model_fn, fourier, n_physical)

Create a model function wrapper that handles Fourier conversion.

The wrapper converts Fourier coefficients to per-angle values before calling the underlying model function.

Parameters:

• model_fn (Callable[[ndarray, ndarray], ndarray]) – Original model function that expects per-angle parameters: f(params, x) where params = [contrast_per_angle, offset_per_angle, physical]

• fourier (FourierReparameterizer) – Fourier reparameterizer instance.

• n_physical (int) – Number of physical parameters.

Returns:

Wrapped model function that accepts Fourier parameters: f(params, x) where params = [fourier_coeffs, physical]

Return type:

Callable[[ndarray, ndarray], ndarray]

Key Classes

homodyne.optimization.nlsq.fourier_reparam.FourierReparamConfig	Configuration for Fourier reparameterization.
homodyne.optimization.nlsq.fourier_reparam.FourierReparameterizer	Handles conversion between Fourier coefficients and per-angle values.

Hierarchical Optimization (Layer 2)

Alternates between physical and per-angle parameter optimization to break gradient cancellation.

Hierarchical Two-Stage Optimization for Anti-Degeneracy Defense.

This module implements alternating optimization between physical and per-angle parameters, breaking the gradient cancellation cycle that causes structural degeneracy in streaming optimization.

Part of Anti-Degeneracy Defense System v2.9.0. See: docs/specs/anti-degeneracy-defense-v2.9.0.md

Algorithm:

Initialize: params = [per_angle_params, physical_params]

for outer_iter in range(max_outer_iterations):

    # Stage 1: Fit PHYSICAL params only
    freeze(per_angle_params)
    result1 = L-BFGS(
        loss_fn(physical_params | frozen_per_angle),
        physical_params
    )
    physical_params = result1.x

    # Stage 2: Fit PER-ANGLE params only
    freeze(physical_params)
    result2 = L-BFGS(
        loss_fn(per_angle_params | frozen_physical),
        per_angle_params
    )
    per_angle_params = result2.x

    # Check convergence
    if converged(physical_params, previous_physical_params):
        break

return [per_angle_params, physical_params]

Why It Works

1. In Stage 1, there are NO per-angle DoF to compete with physical params

2. gamma_dot_t0 gradient CANNOT cancel (no per-angle params to absorb signal)

3. Physical params converge to true values

4. Stage 2 only cleans up residuals with physical interpretation fixed

class homodyne.optimization.nlsq.hierarchical.HierarchicalConfig

Bases: object

Configuration for hierarchical optimization.

enable

Whether to enable hierarchical optimization. Default True.

Type:

bool

max_outer_iterations

Maximum outer iterations. Default 5.

Type:

int

outer_tolerance

Convergence tolerance for physical parameters. Default 1e-6.

Type:

float

physical_max_iterations

Max iterations for Stage 1 (physical params). Default 100.

Type:

int

physical_ftol

Function tolerance for Stage 1. Default 1e-8.

Type:

float

per_angle_max_iterations

Max iterations for Stage 2 (per-angle params). Default 50.

Type:

int

per_angle_ftol

Function tolerance for Stage 2. Default 1e-6.

Type:

float

log_stage_transitions

Whether to log stage transitions. Default True.

Type:

bool

save_intermediate_results

Whether to save intermediate results. Default False.

Type:

bool

enable: bool = True

max_outer_iterations: int = 5

outer_tolerance: float = 1e-06

physical_max_iterations: int = 100

physical_ftol: float = 1e-08

per_angle_max_iterations: int = 50

per_angle_ftol: float = 1e-06

log_stage_transitions: bool = True

save_intermediate_results: bool = False

classmethod from_dict(config_dict)

Create config from dictionary with safe type conversion.

Return type:

HierarchicalConfig

__init__(enable=True, max_outer_iterations=5, outer_tolerance=1e-06, physical_max_iterations=100, physical_ftol=1e-08, per_angle_max_iterations=50, per_angle_ftol=1e-06, log_stage_transitions=True, save_intermediate_results=False)

class homodyne.optimization.nlsq.hierarchical.HierarchicalResult

Bases: object

Result from hierarchical optimization.

x

Optimized parameters.

Type:

np.ndarray

fun

Final loss value.

Type:

float

success

Whether optimization succeeded.

Type:

bool

n_outer_iterations

Number of outer iterations performed.

Type:

int

history

History of each outer iteration.

Type:

list

total_time

Total optimization time in seconds.

Type:

float

message

Status message.

Type:

str

x: ndarray

fun: float

success: bool

n_outer_iterations: int

history: list[dict]

total_time: float = 0.0

message: str = ''

__init__(x, fun, success, n_outer_iterations, history=<factory>, total_time=0.0, message='')

class homodyne.optimization.nlsq.hierarchical.HierarchicalOptimizer

Bases: object

Two-stage hierarchical optimizer for decoupled fitting.

This optimizer breaks the gradient cancellation problem by alternating between physical and per-angle parameter optimization:

Stage 1: Physical parameters only

• Per-angle parameters are frozen

• gamma_dot_t0 gradient cannot be cancelled by per-angle absorption

• Physical params converge to true values

Stage 2: Per-angle parameters only

• Physical parameters are frozen

• Per-angle params absorb only experimental noise

• Cannot change the physical interpretation

Parameters:

• config (HierarchicalConfig) – Hierarchical optimization configuration.

• n_phi (int) – Number of unique phi angles.

• n_physical (int) – Number of physical parameters.

• fourier_reparameterizer (FourierReparameterizer | None) – Fourier reparameterizer if using Fourier mode.

Examples

>>> config = HierarchicalConfig(max_outer_iterations=5)
>>> optimizer = HierarchicalOptimizer(config, n_phi=23, n_physical=7)
>>> result = optimizer.fit(loss_fn, grad_fn, p0, bounds)

__init__(config, n_phi, n_physical, fourier_reparameterizer=None)

Initialize hierarchical optimizer.

Parameters:

• config (HierarchicalConfig) – Configuration.

• n_phi (int) – Number of unique phi angles.

• n_physical (int) – Number of physical parameters.

• fourier_reparameterizer (FourierReparameterizer | None) – Fourier reparameterizer for Fourier mode.

per_angle_indices: ndarray

physical_indices: ndarray

fit(loss_fn, grad_fn, p0, bounds, outer_iteration_callback=None)

Run hierarchical optimization.

Parameters:

• loss_fn (Callable[[ndarray], float]) – Loss function f(params) -> scalar.

• grad_fn (Callable[[ndarray], ndarray] | None) – Gradient function g(params) -> gradient array. If None, uses finite differences.

• p0 (ndarray) – Initial parameters.

• bounds (tuple[ndarray, ndarray]) – (lower_bounds, upper_bounds).

• outer_iteration_callback (Callable[[ndarray, int], None] | None) – Optional callback called at the start of each outer iteration. Signature: callback(current_params, outer_iter). Used for updating shear-sensitivity weights based on current phi0 estimate.

Returns:

Optimization result with diagnostics.

Return type:

HierarchicalResult

get_diagnostics()

Get optimizer diagnostics.

Returns:

Diagnostic information.

Return type:

dict

Key Classes

homodyne.optimization.nlsq.hierarchical.HierarchicalConfig	Configuration for hierarchical optimization.
homodyne.optimization.nlsq.hierarchical.HierarchicalResult	Result from hierarchical optimization.
homodyne.optimization.nlsq.hierarchical.HierarchicalOptimizer	Two-stage hierarchical optimizer for decoupled fitting.

Adaptive Regularization (Layer 3)

CV-based regularization with automatic lambda tuning.

Adaptive Relative Regularization for Anti-Degeneracy Defense.

This module implements CV-based (Coefficient of Variation) regularization that scales properly with data, replacing the ineffective absolute variance regularization.

Part of Anti-Degeneracy Defense System v2.9.0. See: docs/specs/anti-degeneracy-defense-v2.9.0.md

Mathematical Formulation:

Current (ineffective):
    L_reg = lambda * Var(params) * n_points

Proposed (CV-based):
    CV = std(params) / abs(mean(params))
    L_reg = lambda * CV^2 * MSE * n_points

Auto-tuned lambda:
    lambda = target_contribution / target_cv^2

    Example: Allow 10% variation (CV=0.1), contribute 10% to loss
    lambda = 0.1 / 0.01 = 10

class homodyne.optimization.nlsq.adaptive_regularization.AdaptiveRegularizationConfig

Bases: object

Configuration for adaptive relative regularization.

enable

Whether to enable regularization. Default True.

Type:

bool

mode

Regularization mode: “absolute”, “relative”, or “auto”. - “absolute”: Original variance-based (L_reg = λ × Var × n) - “relative”: CV-based (L_reg = λ × CV² × MSE × n) - “auto”: Use relative for n_phi > 5, absolute otherwise

Type:

str

lambda_base

Base regularization strength. Default 1.0 (100× stronger than v2.8).

Type:

float

target_cv

Target coefficient of variation. Default 0.10 (10% variation allowed).

Type:

float

target_contribution

Target fraction of MSE to contribute. Default 0.10 (10% of loss).

Type:

float

auto_tune_lambda

Whether to auto-compute λ from target_cv and target_contribution.

Type:

bool

max_cv

Maximum allowed CV before hard constraint warning. Default 0.20.

Type:

float

group_indices

Parameter group indices [(start, end), …]. Auto-computed if None.

Type:

list of tuple, optional

enable: bool = True

mode: Literal['absolute', 'relative', 'auto'] = 'relative'

lambda_base: float = 1.0

target_cv: float = 0.1

target_contribution: float = 0.1

auto_tune_lambda: bool = True

max_cv: float = 0.2

group_indices: list[tuple[int, int]] | None = None

classmethod from_dict(config_dict)

Create config from dictionary with safe type conversion.

Return type:

AdaptiveRegularizationConfig

__init__(enable=True, mode='relative', lambda_base=1.0, target_cv=0.1, target_contribution=0.1, auto_tune_lambda=True, max_cv=0.2, group_indices=None)

class homodyne.optimization.nlsq.adaptive_regularization.AdaptiveRegularizer

Bases: object

CV-based adaptive regularization for per-angle parameters.

This regularizer addresses the fundamental problem where absolute variance regularization (λ=0.01) contributed only ~0.05% to total loss, providing no effective constraint on per-angle parameter variation.

The CV-based approach ensures regularization scales properly: - CV is dimensionless (ratio of std to mean) - Auto-tuned λ makes regularization ~10% of MSE - Prevents per-angle parameters from absorbing physical signals

Parameters:

• config (AdaptiveRegularizationConfig) – Regularization configuration.

• n_phi (int) – Number of unique phi angles.

lambda_value

Effective regularization strength (auto-tuned or from config).

Type:

float

group_indices

Parameter groups to regularize.

Type:

list of tuple

Examples

>>> config = AdaptiveRegularizationConfig(target_cv=0.10, target_contribution=0.10)
>>> regularizer = AdaptiveRegularizer(config, n_phi=23)
>>> reg_term = regularizer.compute_regularization(
...     params, mse=0.04, n_points=23_000_000
... )

__init__(config, n_phi, n_params=None)

Initialize adaptive regularizer.

Parameters:

• config (AdaptiveRegularizationConfig) – Regularization configuration.

• n_phi (int) – Number of unique phi angles.

• n_params (int | None) – Actual parameter vector length. When provided and less than 2 * n_phi + n_physical, auto_averaged mode is assumed (2 scaling params instead of 2 * n_phi).

compute_regularization(params, mse, n_points)

Compute regularization term to add to loss.

Parameters:

• params (ndarray) – Full parameter vector.

• mse (float) – Current mean squared error.

• n_points (int) – Number of data points.

Returns:

Regularization term to add to loss (SSE scale).

Return type:

float

compute_regularization_jax(params, mse, n_points)

Compute regularization term using JAX for autodiff compatibility.

This method uses JAX operations (jnp) instead of NumPy, making it compatible with JAX’s JIT compilation and autodiff (jax.grad).

Use this method when the regularization needs to be part of a differentiable loss function.

Parameters:

• params (Array) – Full parameter vector (JAX array, possibly traced).

• mse (Array) – Current mean squared error (JAX scalar, possibly traced).

• n_points (int) – Number of data points.

Returns:

Regularization term to add to loss (SSE scale, JAX scalar).

Return type:

Array

compute_regularization_gradient(params, mse, n_points)

Compute gradient of regularization term.

Parameters:

• params (ndarray) – Full parameter vector.

• mse (float) – Current mean squared error.

• n_points (int) – Number of data points.

Returns:

Gradient w.r.t. all parameters (zeros for non-regularized params).

Return type:

ndarray

check_constraint_violation(params)

Check if CV exceeds max_cv threshold.

Parameters:

params (ndarray) – Full parameter vector.

Returns:

Dictionary of violations, empty if none.

Return type:

dict[str, dict]

get_diagnostics()

Get regularization diagnostics for logging.

Returns:

Diagnostic information including CV values and contribution.

Return type:

dict

log_summary(params, mse, n_points)

Log regularization summary.

Parameters:

• params (ndarray) – Full parameter vector.

• mse (float) – Current mean squared error.

• n_points (int) – Number of data points.

Return type:

None

Key Classes

homodyne.optimization.nlsq.adaptive_regularization.AdaptiveRegularizationConfig	Configuration for adaptive relative regularization.
homodyne.optimization.nlsq.adaptive_regularization.AdaptiveRegularizer	CV-based adaptive regularization for per-angle parameters.

Gradient Collapse Monitor (Layer 4)

Runtime detection of gradient collapse with automatic response actions.

Gradient Collapse Monitor for Anti-Degeneracy Defense.

This module provides runtime detection of gradient collapse (physical params losing gradient signal) with automatic response actions.

Part of Anti-Degeneracy Defense System v2.9.0. See: docs/specs/anti-degeneracy-defense-v2.9.0.md

Detection Mechanism:

Monitor the ratio:
    ratio = norm(grad_physical) / norm(grad_per_angle)

If ratio < threshold for N consecutive iterations:
    - Gradient collapse detected
    - Physical params are losing signal to per-angle params

Response Actions

• “warn”: Log warning only

• “hierarchical”: Switch to hierarchical optimization mode

• “reset”: Reset per-angle params to mean values

• “abort”: Abort optimization and return best params so far

class homodyne.optimization.nlsq.gradient_monitor.GradientMonitorConfig

Bases: object

Configuration for gradient collapse detection.

enable

Whether to enable gradient monitoring. Default True.

Type:

bool

ratio_threshold

Ratio of norm(grad_physical) / norm(grad_per_angle) below this triggers detection. Default 0.01 (physical gradient is 1% of per-angle gradient).

Type:

float

consecutive_triggers

Must trigger N consecutive times to confirm collapse. Default 5.

Type:

int

response_mode

Response action on collapse detection: - “warn”: Log warning only - “hierarchical”: Switch to hierarchical optimization - “reset”: Reset per-angle params to mean - “abort”: Abort and return best params

Type:

str

reset_per_angle_to_mean

When resetting, reset per-angle to mean values. Default True.

Type:

bool

lambda_multiplier_on_collapse

Multiply regularization λ by this on collapse. Default 10.0.

Type:

float

check_interval

Check every N iterations. Default 1 (every iteration).

Type:

int

enable: bool = True

ratio_threshold: float = 0.01

consecutive_triggers: int = 5

response_mode: Literal['warn', 'hierarchical', 'reset', 'abort'] = 'hierarchical'

reset_per_angle_to_mean: bool = True

lambda_multiplier_on_collapse: float = 10.0

check_interval: int = 1

watch_parameters: list[int] | None = None

watch_threshold: float = 1e-08

watch_consecutive_triggers: int = 3

watch_min_iteration: int = 5

classmethod from_dict(config_dict)

Create config from dictionary with safe type conversion.

Return type:

GradientMonitorConfig

__init__(enable=True, ratio_threshold=0.01, consecutive_triggers=5, response_mode='hierarchical', reset_per_angle_to_mean=True, lambda_multiplier_on_collapse=10.0, check_interval=1, watch_parameters=None, watch_threshold=1e-08, watch_consecutive_triggers=3, watch_min_iteration=5)

class homodyne.optimization.nlsq.gradient_monitor.CollapseEvent

Bases: object

Record of a gradient collapse event.

iteration

Iteration when collapse was detected.

Type:

int

ratio

Gradient ratio at detection.

Type:

float

physical_grad_norm

Physical parameter gradient norm.

Type:

float

per_angle_grad_norm

Per-angle parameter gradient norm.

Type:

float

response_mode

Response action taken.

Type:

str

iteration: int

ratio: float

physical_grad_norm: float

per_angle_grad_norm: float

response_mode: str

__init__(iteration, ratio, physical_grad_norm, per_angle_grad_norm, response_mode)

class homodyne.optimization.nlsq.gradient_monitor.GradientCollapseMonitor

Bases: object

Monitor for detecting and responding to gradient collapse.

This monitor tracks the ratio of physical to per-angle gradient norms during optimization. When the ratio drops below a threshold for consecutive iterations, it indicates that physical parameters are losing gradient signal (being absorbed by per-angle parameters).

Parameters:

• config (GradientMonitorConfig) – Monitor configuration.

• physical_indices (Sequence[int] | ndarray) – Indices of physical parameters in the full parameter vector.

• per_angle_indices (Sequence[int] | ndarray) – Indices of per-angle parameters in the full parameter vector.

collapse_detected

Whether gradient collapse has been detected.

Type:

bool

consecutive_count

Current count of consecutive low-ratio iterations.

Type:

int

Notes

History is capped at MAX_HISTORY_SIZE to prevent memory leaks during long-running optimizations. Older entries are discarded when the limit is reached.

Examples

>>> config = GradientMonitorConfig(ratio_threshold=0.01, consecutive_triggers=5)
>>> monitor = GradientCollapseMonitor(config, physical_indices=[6,7,8,9,10,11,12],
...                                    per_angle_indices=list(range(6)))
>>> for iter in range(100):
...     gradients = compute_gradients(params)
...     status = monitor.check(gradients, iter)
...     if status == "COLLAPSE_DETECTED":
...         response = monitor.get_response()
...         # Take action based on response

MAX_HISTORY_SIZE: int = 1000

__init__(config, physical_indices, per_angle_indices)

Initialize gradient collapse monitor.

Parameters:

• config (GradientMonitorConfig) – Monitor configuration.

• physical_indices (Sequence[int] | ndarray) – Indices of physical parameters. Converted to numpy array internally to support both NumPy and JAX array indexing.

• per_angle_indices (Sequence[int] | ndarray) – Indices of per-angle parameters (or Fourier coefficients when Fourier reparameterization is active). Converted to numpy array internally.

Notes

When Fourier reparameterization is active, per_angle_indices should correspond to Fourier coefficient indices (typically 10 for order=2), not independent per-angle indices (2 * n_phi).

physical_indices: ndarray

per_angle_indices: ndarray

history: deque[dict]

consecutive_count: int

collapse_detected: bool

collapse_events: list[CollapseEvent]

best_params: ndarray | None

best_loss: float

check(gradients, iteration, params=None, loss=None)

Check for gradient collapse.

Parameters:

• gradients (ndarray) – Full gradient vector.

• iteration (int) – Current iteration number.

• params (ndarray | None) – Current parameters (for response actions and tracking).

• loss (float | None) – Current loss value (for tracking best params).

Returns:

Status: “OK”, “WARNING”, “COLLAPSE_DETECTED”

Return type:

str

get_response()

Get response action after collapse detection.

Returns:

Response action dictionary, or None if no collapse.

Return type:

dict | None

compute_reset_params(params, n_phi)

Compute parameters with per-angle values reset to mean.

Parameters:

• params (ndarray) – Current parameter vector.

• n_phi (int) – Number of phi angles.

Returns:

Parameters with per-angle values reset.

Return type:

ndarray

reset()

Reset monitor state for new optimization run.

Return type:

None

get_diagnostics()

Get monitoring diagnostics for logging.

Returns:

Diagnostic information.

Return type:

dict

log_summary()

Log monitoring summary.

Return type:

None

homodyne.optimization.nlsq.gradient_monitor.create_gradient_function_with_monitoring(grad_fn, monitor)

Wrap gradient function to include monitoring.

Parameters:

• grad_fn (Callable[[ndarray], ndarray]) – Original gradient function.

• monitor (GradientCollapseMonitor) – Monitor instance.

Returns:

Wrapped gradient function that records to monitor.

Return type:

Callable[[ndarray], ndarray]

Key Classes

homodyne.optimization.nlsq.gradient_monitor.GradientMonitorConfig	Configuration for gradient collapse detection.
homodyne.optimization.nlsq.gradient_monitor.CollapseEvent	Record of a gradient collapse event.
homodyne.optimization.nlsq.gradient_monitor.GradientCollapseMonitor	Monitor for detecting and responding to gradient collapse.

Shear-Sensitivity Weighting (Layer 5)

Weights residuals by |cos(φ₀-φ)| to prevent gradient cancellation. Computed in Homodyne and passed to NLSQ as generic residual weights.

Shear-Sensitivity Weighting for Anti-Degeneracy Defense.

This module implements angle-dependent loss weighting to prevent gradient cancellation in the shear term during optimization.

Part of Anti-Degeneracy Defense System v2.9.1.

The Problem

The shear term gradient is:

d(g1_shear)/d(gamma_dot_t0) ~ cos(phi0 - phi)

When summed uniformly over all angles: - Angles near phi0: cos(phi0 - phi) ~ +1 (positive contribution) - Angles near phi0 +/- 90deg: cos ~ 0 (negligible) - Angles near phi0 +/- 180deg: cos ~ -1 (negative contribution)

With uniformly distributed angles, positive and negative contributions CANCEL, leading to near-zero net gradient for gamma_dot_t0. This causes the shear parameter to collapse to its lower bound.

The Solution

Use angle-dependent loss weighting:

L = sum_phi w(phi) * sum_tau (g2_model - g2_exp)^2

where w(phi) emphasizes shear-sensitive angles:

w(phi) = w_min + (1 - w_min) * abs(cos(phi0_current - phi))^alpha

This converts gradient cancellation into a weighted sum where shear-sensitive angles (parallel/antiparallel to flow) contribute more than perpendicular angles. All angles still contribute to prevent information loss.

Configuration

shear_weighting:

enable: true # Enable shear-sensitivity weighting min_weight: 0.3 # Minimum weight (0-1) alpha: 1.0 # Shear sensitivity exponent (1 = linear) update_frequency: 1 # Update weights every N outer iterations initial_phi0: null # Initial phi0 guess (null = use config)

class homodyne.optimization.nlsq.shear_weighting.ShearWeightingConfig

Bases: object

Configuration for shear-sensitivity weighting.

enable

Enable shear-sensitivity weighting. Default True.

Type:

bool

min_weight

Minimum weight for perpendicular angles. Range [0, 1]. Default 0.3.

Type:

float

alpha

Shear sensitivity exponent. Higher = more aggressive weighting. Default 1.0 (linear).

Type:

float

update_frequency

Update weights every N outer iterations. Default 1.

Type:

int

initial_phi0

Initial phi0 guess in degrees. None = use config or 0.0.

Type:

float or None

normalize

Normalize weights so sum = n_phi. Default True.

Type:

bool

enable: bool = True

min_weight: float = 0.3

alpha: float = 1.0

update_frequency: int = 1

initial_phi0: float | None = None

normalize: bool = True

classmethod from_config(config)

Create from configuration dictionary.

Parameters:

config (Mapping) – Configuration dictionary.

Returns:

Configuration object.

Return type:

ShearWeightingConfig

__init__(enable=True, min_weight=0.3, alpha=1.0, update_frequency=1, initial_phi0=None, normalize=True)

class homodyne.optimization.nlsq.shear_weighting.ShearSensitivityWeighting

Bases: object

Shear-sensitivity weighted loss for anti-degeneracy defense.

This class manages angle-dependent weights that emphasize shear-sensitive angles during optimization, preventing gradient cancellation.

Parameters:

• phi_angles (ndarray) – Array of phi angles in degrees.

• n_physical (int) – Number of physical parameters.

• phi0_index (int) – Index of phi0 in physical parameters (typically 6 for laminar_flow).

• config (ShearWeightingConfig | None) – Weighting configuration.

Examples

>>> phi_angles = np.array([-30, 0, 30, 60, 90, 120])
>>> weighter = ShearSensitivityWeighting(phi_angles, n_physical=7, phi0_index=6)
>>> weights = weighter.get_weights(phi0_current=-5.0)
>>> # Angles near -5 deg and 175 deg get higher weight

__init__(phi_angles, n_physical, phi0_index, config=None)

update_phi0(params, iteration=0)

Update phi0 estimate from current parameters.

Parameters:

• params (ndarray) – Current parameter vector. Physical parameters should be at the end.

• iteration (int) – Current iteration number.

Return type:

None

get_weights(phi0_current=None)

Get current angle weights.

Parameters:

phi0_current (float | None) – Override phi0 for weight computation. If None, uses stored value.

Returns:

Weight array of shape (n_phi,).

Return type:

ndarray

get_weights_jax()

Get current angle weights as JAX array.

Returns:

Weight array of shape (n_phi,).

Return type:

Array

apply_weights_to_loss(residuals, phi_indices)

Apply angle weights to residuals for loss computation.

Computes weighted mean squared error:

L = sum_i w[phi_idx[i]] * residuals[i]^2 / sum_i w[phi_idx[i]]

Parameters:

• residuals (Array) – Residuals array of shape (n_data,).

• phi_indices (Array) – Phi index for each data point, shape (n_data,).

Returns:

Weighted loss (scalar).

Return type:

Array

compute_weighted_mse(residuals, phi_indices)

Compute weighted MSE (for gradient computation).

Parameters:

• residuals (Array) – Residuals array of shape (n_data,).

• phi_indices (Array) – Phi index for each data point, shape (n_data,).

Returns:

Weighted MSE (scalar).

Return type:

Array

get_diagnostics()

Get weighting diagnostics.

Returns:

Diagnostic information.

Return type:

dict

property phi0_current: float

Current phi0 estimate in degrees.

homodyne.optimization.nlsq.shear_weighting.create_shear_weighting(phi_angles, n_physical, config=None, physical_param_names=None)

Factory function to create shear weighting if enabled.

Parameters:

• phi_angles (ndarray) – Phi angles in degrees.

• n_physical (int) – Number of physical parameters.

• config (Mapping | None) – Configuration dictionary.

Returns:

Weighting object if enabled, None otherwise.

Return type:

ShearSensitivityWeighting | None

Key Classes

homodyne.optimization.nlsq.shear_weighting.ShearWeightingConfig	Configuration for shear-sensitivity weighting.
homodyne.optimization.nlsq.shear_weighting.ShearSensitivityWeighting	Shear-sensitivity weighted loss for anti-degeneracy defense.

Anti-Degeneracy Controller

Unified controller that orchestrates all defense layers.

Anti-Degeneracy Controller - Orchestrator for 5-Layer Defense System.

This module provides a clean interface for initializing and coordinating the 5-layer anti-degeneracy defense system for NLSQ optimization.

The controller encapsulates: - Layer 1: Fourier/Constant Reparameterization - Layer 2: Hierarchical Optimization - Layer 3: Adaptive CV-based Regularization - Layer 4: Gradient Collapse Monitoring - Layer 5: Shear-Sensitivity Weighting

Usage:

controller = AntiDegeneracyController.from_config(
    config_dict, n_phi, phi_angles, n_physical
)
if controller.is_enabled:
    # Use controller.fourier, controller.hierarchical, etc.
    transformed_params = controller.transform_params_to_fourier(initial_params)
    model_fn = controller.wrap_model_fn(base_model_fn)

Version: 2.9.0 Author: Claude Code

class homodyne.optimization.nlsq.anti_degeneracy_controller.AntiDegeneracyConfig

Bases: object

Configuration for the Anti-Degeneracy Defense System.

enable

Master switch for all anti-degeneracy defenses.

Type:

bool

per_angle_mode

Mode for per-angle parameters: “individual”, “constant”, “fourier”, or “auto”.

Type:

str

fourier_order

Order of Fourier series (order=2 -> 5 coefficients per group).

Type:

int

fourier_auto_threshold

n_phi threshold for auto mode to switch to Fourier.

Type:

int

constant_scaling_threshold

n_phi threshold for auto mode to use constant scaling (n_phi >= threshold).

Type:

int

hierarchical_enable

Enable hierarchical two-stage optimization.

Type:

bool

hierarchical_max_outer_iterations

Maximum outer iterations for hierarchical optimization.

Type:

int

hierarchical_outer_tolerance

Convergence tolerance on physical parameter change.

Type:

float

regularization_mode

Regularization mode: “absolute”, “relative”, or “auto”.

Type:

str

regularization_lambda

Base regularization strength.

Type:

float

regularization_target_cv

Target coefficient of variation (0-1).

Type:

float

regularization_target_contribution

Target regularization contribution to loss (0-1).

Type:

float

gradient_monitoring_enable

Enable gradient collapse monitoring.

Type:

bool

gradient_ratio_threshold

Collapse threshold for norm(grad_physical)/norm(grad_per_angle).

Type:

float

gradient_consecutive_triggers

Number of consecutive triggers to confirm collapse.

Type:

int

gradient_response_mode

Response action: “warn”, “hierarchical”, “reset”, “abort”.

Type:

str

enable: bool = True

per_angle_mode: str = 'auto'

fourier_order: int = 2

fourier_auto_threshold: int = 6

constant_scaling_threshold: int = 3

hierarchical_enable: bool = True

hierarchical_max_outer_iterations: int = 5

hierarchical_outer_tolerance: float = 1e-06

hierarchical_physical_max_iterations: int = 100

hierarchical_per_angle_max_iterations: int = 50

regularization_mode: str = 'relative'

regularization_lambda: float = 1.0

regularization_target_cv: float = 0.1

regularization_target_contribution: float = 0.1

regularization_max_cv: float = 0.2

gradient_monitoring_enable: bool = True

gradient_ratio_threshold: float = 0.01

gradient_consecutive_triggers: int = 5

gradient_response_mode: str = 'hierarchical'

shear_weighting_enable: bool = True

shear_weighting_min_weight: float = 0.3

shear_weighting_alpha: float = 1.0

shear_weighting_update_frequency: int = 1

shear_weighting_normalize: bool = True

classmethod from_dict(config_dict)

Create config from nested dictionary.

Parameters:

config_dict (dict[str, Any]) –

Configuration dictionary with structure:

{
    "enable": bool,
    "per_angle_mode": str,
    "fourier_order": int,
    "fourier_auto_threshold": int,
    "hierarchical": {...},
    "regularization": {...},
    "gradient_monitoring": {...}
}

Returns:

Validated configuration object.

Return type:

AntiDegeneracyConfig

__init__(enable=True, per_angle_mode='auto', fourier_order=2, fourier_auto_threshold=6, constant_scaling_threshold=3, hierarchical_enable=True, hierarchical_max_outer_iterations=5, hierarchical_outer_tolerance=1e-06, hierarchical_physical_max_iterations=100, hierarchical_per_angle_max_iterations=50, regularization_mode='relative', regularization_lambda=1.0, regularization_target_cv=0.1, regularization_target_contribution=0.1, regularization_max_cv=0.2, gradient_monitoring_enable=True, gradient_ratio_threshold=0.01, gradient_consecutive_triggers=5, gradient_response_mode='hierarchical', shear_weighting_enable=True, shear_weighting_min_weight=0.3, shear_weighting_alpha=1.0, shear_weighting_update_frequency=1, shear_weighting_normalize=True)

class homodyne.optimization.nlsq.anti_degeneracy_controller.AntiDegeneracyController

Bases: object

Orchestrator for the 5-Layer Anti-Degeneracy Defense System.

This controller provides a clean interface for initializing and coordinating all anti-degeneracy components.

config

Configuration for the defense system.

Type:

AntiDegeneracyConfig

n_phi

Number of phi angles.

Type:

int

n_physical

Number of physical parameters.

Type:

int

phi_angles

Array of phi angles in radians.

Type:

np.ndarray

fourier

Layer 1: Fourier reparameterization component.

Type:

FourierReparameterizer | None

hierarchical

Layer 2: Hierarchical optimization component.

Type:

HierarchicalOptimizer | None

regularizer

Layer 3: Adaptive regularization component.

Type:

AdaptiveRegularizer | None

monitor

Layer 4: Gradient collapse monitoring component.

Type:

GradientCollapseMonitor | None

shear_weighter

Layer 5: Shear-sensitivity weighting component.

Type:

ShearSensitivityWeighting | None

per_angle_mode_actual

Actual mode used (“constant”, “fourier”, or “independent”).

Type:

str

config: AntiDegeneracyConfig

n_phi: int

n_physical: int

phi_angles: ndarray

fourier: FourierReparameterizer | None = None

hierarchical: HierarchicalOptimizer | None = None

regularizer: AdaptiveRegularizer | None = None

monitor: GradientCollapseMonitor | None = None

shear_weighter: ShearSensitivityWeighting | None = None

mapper: ParameterIndexMapper | None = None

per_angle_mode_actual: str = 'disabled'

classmethod from_config(config_dict, n_phi, phi_angles, n_physical, per_angle_scaling=True, is_laminar_flow=True)

Create controller from configuration dictionary.

Parameters:

• config_dict (dict[str, Any]) – Anti-degeneracy configuration dictionary.

• n_phi (int) – Number of phi angles.

• phi_angles (ndarray) – Array of phi angles in radians.

• n_physical (int) – Number of physical parameters (7 for laminar_flow).

• per_angle_scaling (bool) – Whether per-angle scaling is enabled.

• is_laminar_flow (bool) – Whether this is laminar_flow mode.

Returns:

Initialized controller with all components.

Return type:

AntiDegeneracyController

property is_enabled: bool

Check if the defense system is enabled and initialized.

property use_fourier: bool

Check if Fourier reparameterization is active.

property use_constant: bool

Check if constant scaling mode is active (either auto_averaged or fixed_constant).

Both modes use constant-style parameter mapping (9 params for auto_averaged, 7 params for fixed_constant), as opposed to individual mode (7 + 2*n_phi params).

property use_fixed_scaling: bool

Check if using FIXED per-angle scaling (7 params, not optimized).

Returns True only for explicit constant mode (“fixed_constant”), where: - Per-angle contrast/offset are FIXED from quantile estimation - Only 7 physical parameters are optimized - Scaling is NOT part of the optimization

This is DIFFERENT from auto_averaged mode, where: - Averaged contrast/offset ARE optimized (9 params total)

property use_averaged_scaling: bool

Check if using OPTIMIZED averaged scaling (9 params).

Returns True only for auto mode with n_phi >= threshold (“auto_averaged”), where: - N per-angle quantile estimates are averaged to 1 contrast + 1 offset - These 2 averaged values ARE OPTIMIZED along with 7 physical params - Total: 9 parameters

property use_hierarchical: bool

Check if hierarchical optimization is active.

property use_shear_weighting: bool

Check if shear-sensitivity weighting is active.

property n_per_angle_params: int

Get the number of per-angle parameters (optimized scaling params).

Returns: - fixed_constant: 0 (scaling is FIXED, not optimized) - auto_averaged: 2 (one contrast, one offset - OPTIMIZED) - fourier: n_coeffs (Fourier coefficients - OPTIMIZED) - individual: 2*n_phi (per-angle contrast + offset - OPTIMIZED)

transform_params_to_fourier(params)

Transform per-angle parameters to Fourier coefficients.

Parameters:

params (ndarray) – Full parameter array: [contrast(n_phi), offset(n_phi), physical].

Returns:

(fourier_params, original_bounds_if_transformed) fourier_params: [contrast_coeffs, offset_coeffs, physical] bounds: (lower, upper) in Fourier space if transformation applied

Return type:

tuple[ndarray, tuple[ndarray, ndarray] | None]

transform_params_from_fourier(fourier_params)

Transform Fourier coefficients back to per-angle parameters.

Parameters:

fourier_params (ndarray) – Fourier parameter array: [contrast_coeffs, offset_coeffs, physical].

Returns:

Per-angle parameter array: [contrast(n_phi), offset(n_phi), physical].

Return type:

ndarray

transform_params_to_constant(params)

Transform per-angle parameters to constant mode.

Computes mean contrast and offset across all angles.

Parameters:

params (ndarray) – Full parameter array: [contrast(n_phi), offset(n_phi), physical].

Returns:

Constant mode parameters: [contrast_mean, offset_mean, physical].

Return type:

ndarray

transform_params_from_constant(constant_params)

Transform constant mode parameters to per-angle form.

Expands single contrast/offset values to all angles.

Parameters:

constant_params (ndarray) – Constant mode parameters: [contrast, offset, physical].

Returns:

Per-angle parameters: [contrast(n_phi), offset(n_phi), physical].

Return type:

ndarray

get_group_variance_indices()

Get group variance indices for NLSQ regularization.

T024: Delegates to ParameterIndexMapper for consistent index calculation regardless of Fourier mode.

Returns:

List of (start, end) tuples for each parameter group.

Return type:

list[tuple[int, int]] | None

get_diagnostics()

Get comprehensive diagnostics from all components.

Returns:

Nested diagnostics from all 5 layers.

Return type:

dict[str, Any]

reset_monitor()

Reset the gradient collapse monitor state.

Return type:

None

get_shear_weights()

Get shear-sensitivity weights for residuals.

Returns:

Array of weights (one per phi angle), or None if not enabled.

Return type:

ndarray | None

update_shear_phi0(params, iteration=0)

Update the phi0 value in shear weighter.

Parameters:

• params (ndarray) – Current parameter vector.

• iteration (int) – Current iteration number.

Return type:

None

compute_fixed_per_angle_scaling(stratified_data, contrast_bounds=(0.0, 1.0), offset_bounds=(0.5, 1.5))

Compute and store fixed per-angle contrast/offset from quantiles.

This method uses physics-informed quantile analysis to estimate contrast and offset for each phi angle independently.

In “constant” mode (v2.17.0+): 1. Computes N contrast + N offset values from quantile estimation 2. These are averaged to 1 contrast + 1 offset for optimization 3. Optimizer works with 9 parameters: 7 physical + 2 averaged scaling 4. The individual per-angle estimates are stored for diagnostics

Parameters:

• stratified_data (Any) – Data containing per-angle g2_flat, phi_flat, t1_flat, t2_flat arrays.

• contrast_bounds (tuple[float, float]) – Valid bounds for contrast parameter.

• offset_bounds (tuple[float, float]) – Valid bounds for offset parameter.

Return type:

None

Notes

This method should be called before optimization when using per_angle_mode=”constant”. The estimates can be retrieved using get_fixed_per_angle_scaling().

Unlike least-squares estimation, this approach: 1. Does not require a model (purely data-driven) 2. Uses physics-informed quantile analysis 3. Is robust to outliers

get_fixed_per_angle_scaling()

Get the fixed per-angle contrast/offset estimates.

Returns:

(contrast_per_angle, offset_per_angle) if computed, None otherwise.

Return type:

tuple[ndarray, ndarray] | None

has_fixed_per_angle_scaling()

Check if fixed per-angle scaling has been computed.

Returns:

True if fixed scaling is available.

Return type:

bool

create_nlsq_callbacks()

Create callbacks for NLSQ’s CurveFit integration.

This method creates callbacks that can be passed to NLSQ’s CurveFit or AdaptiveHybridStreamingOptimizer to enable anti-degeneracy defenses.

Returns:

Dictionary of callbacks compatible with NLSQ: - ‘loss_augmentation’: Callable for regularization loss - ‘iteration_callback’: Callable for gradient monitoring - ‘group_variance_indices’: Indices for group variance regularization

Return type:

dict[str, Any]

Notes

For NLSQ v0.4+, callbacks can be passed to CurveFit.curve_fit() or injected into HybridStreamingConfig.

Example

>>> controller = AntiDegeneracyController.from_config(config, n_phi, phi_angles, n_physical)
>>> callbacks = controller.create_nlsq_callbacks()
>>> result = fitter.curve_fit(f, xdata, ydata, **callbacks)

create_hybrid_streaming_config_kwargs()

Create kwargs for NLSQ’s HybridStreamingConfig.

Returns kwargs that can be used to configure NLSQ’s AdaptiveHybridStreamingOptimizer with anti-degeneracy features.

Returns:

Configuration kwargs for HybridStreamingConfig: - ‘enable_group_variance_regularization’: bool - ‘group_variance_lambda’: float - ‘group_variance_indices’: list[tuple[int, int]]

Return type:

dict[str, Any]

Notes

For NLSQ v0.4+, pass these to HybridStreamingConfig constructor.

Example

>>> controller = AntiDegeneracyController.from_config(...)
>>> kwargs = controller.create_hybrid_streaming_config_kwargs()
>>> config = HybridStreamingConfig(**kwargs)

__init__(config, n_phi, n_physical, phi_angles, fourier=None, hierarchical=None, regularizer=None, monitor=None, shear_weighter=None, mapper=None, per_angle_mode_actual='disabled', _fixed_contrast_per_angle=None, _fixed_offset_per_angle=None, _is_initialized=False)

Key Classes

homodyne.optimization.nlsq.anti_degeneracy_controller.AntiDegeneracyConfig	Configuration for the Anti-Degeneracy Defense System.
homodyne.optimization.nlsq.anti_degeneracy_controller.AntiDegeneracyController	Orchestrator for the 5-Layer Anti-Degeneracy Defense System.

NLSQ Configuration

Configuration dataclasses and utilities for NLSQ optimization.

NLSQ configuration dataclass and validation.

This module provides the NLSQConfig dataclass for parsing and validating NLSQ-specific configuration settings from the YAML config file.

Part of Phase 3 architecture refactoring to reduce wrapper.py complexity.

Config Consolidation (v2.14.0, FR-014): - Single entry point: NLSQConfig.from_yaml() or NLSQConfig.from_dict() - Safe type conversion utilities: safe_float, safe_int - Full validation via validate() method

homodyne.optimization.nlsq.config.safe_float(value, default)

Convert value to float safely, returning default on failure.

Parameters:

• value (Any) – Value to convert to float.

• default (float) – Default value to return if conversion fails.

Returns:

Converted float value or default.

Return type:

float

Examples

>>> safe_float("3.14", 0.0)
3.14
>>> safe_float(None, 1.0)
1.0
>>> safe_float("invalid", 2.5)
2.5

homodyne.optimization.nlsq.config.safe_int(value, default)

Convert value to int safely, returning default on failure.

Parameters:

• value (Any) – Value to convert to int.

• default (int) – Default value to return if conversion fails.

Returns:

Converted int value or default.

Return type:

int

Examples

>>> safe_int("42", 0)
42
>>> safe_int(None, 10)
10
>>> safe_int("invalid", 5)
5

class homodyne.optimization.nlsq.config.HybridRecoveryConfig

Bases: object

Configuration for hybrid streaming optimizer recovery strategy.

T029: Implements 3-attempt recovery with progressively conservative settings.

When the hybrid streaming optimizer fails, it retries with: - Reduced learning rate (0.5× per retry) - Increased regularization (2× per retry) - Smaller trust region (0.5× per retry)

max_retries

Maximum retry attempts. Default: 3.

Type:

int

lr_decay

Learning rate multiplier per retry. Default: 0.5.

Type:

float

lambda_growth

Regularization multiplier per retry. Default: 2.0.

Type:

float

trust_decay

Trust region multiplier per retry. Default: 0.5.

Type:

float

log_retries

Whether to log retry attempts. Default: True.

Type:

bool

max_retries: int = 3

lr_decay: float = 0.5

lambda_growth: float = 2.0

trust_decay: float = 0.5

log_retries: bool = True

get_retry_settings(attempt)

Get settings for a specific retry attempt.

Parameters:

attempt (int) – Retry attempt number (1-based).

Returns:

Settings for this retry attempt.

Return type:

dict

__init__(max_retries=3, lr_decay=0.5, lambda_growth=2.0, trust_decay=0.5, log_retries=True)

class homodyne.optimization.nlsq.config.NLSQConfig

Bases: object

Configuration for NLSQ (Nonlinear Least Squares) optimization.

This dataclass consolidates NLSQ settings that were previously scattered across wrapper.py, improving maintainability and testability.

loss

Loss function for robust fitting. Options: “linear”, “soft_l1”, “huber”, “cauchy”, “arctan”. Default: “soft_l1”.

Type:

str

trust_region_scale

Scale factor for trust region. Default: 1.0.

Type:

float

max_iterations

Maximum number of optimization iterations. Default: 1000.

Type:

int

ftol

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

Type:

float

xtol

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

Type:

float

gtol

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

Type:

float

x_scale

Parameter scaling. “jac” for Jacobian-based, list for manual. Default: “jac”.

Type:

str | list[float] | None

x_scale_map

Per-parameter scaling overrides. Default: None.

Type:

dict[str, float] | None

enable_diagnostics

Whether to compute diagnostics (Jacobian stats, etc.). Default: True.

Type:

bool

enable_streaming

Whether to enable streaming optimizer for large datasets. Default: True.

Type:

bool

streaming_chunk_size

Points per chunk for streaming optimizer. Default: 50000.

Type:

int

enable_stratified

Whether to enable stratified least squares. Default: True.

Type:

bool

target_chunk_size

Target points per chunk for stratified optimization. Default: 100000.

Type:

int

enable_recovery

Whether to enable automatic error recovery. Default: True.

Type:

bool

max_recovery_attempts

Maximum recovery attempts per strategy. Default: 3.

Type:

int

workflow: str = 'auto'

goal: str = 'quality'

loss: str = 'soft_l1'

trust_region_scale: float = 1.0

max_iterations: int = 1000

ftol: float = 1e-08

xtol: float = 1e-08

gtol: float = 1e-08

x_scale: str | list[float] | None = 'jac'

x_scale_map: dict[str, float] | None = None

enable_diagnostics: bool = True

enable_streaming: bool = True

streaming_chunk_size: int = 50000

enable_stratified: bool = True

target_chunk_size: int = 100000

enable_recovery: bool = True

max_recovery_attempts: int = 3

enable_progress_bar: bool = True

verbose: int = 1

log_iteration_interval: int = 10

enable_hybrid_streaming: bool = True

hybrid_normalize: bool = True

hybrid_normalization_strategy: str = 'auto'

hybrid_warmup_iterations: int = 200

hybrid_max_warmup_iterations: int = 500

hybrid_warmup_learning_rate: float = 0.001

hybrid_gauss_newton_max_iterations: int = 100

hybrid_gauss_newton_tol: float = 1e-08

hybrid_chunk_size: int = 10000

hybrid_trust_region_initial: float = 1.0

hybrid_regularization_factor: float = 1e-10

hybrid_enable_checkpoints: bool = True

hybrid_checkpoint_frequency: int = 100

hybrid_validate_numerics: bool = True

hybrid_enable_warm_start_detection: bool = True

hybrid_warm_start_threshold: float = 0.01

hybrid_enable_adaptive_warmup_lr: bool = True

hybrid_warmup_lr_refinement: float = 1e-06

hybrid_warmup_lr_careful: float = 1e-05

hybrid_enable_cost_guard: bool = True

hybrid_cost_increase_tolerance: float = 0.05

hybrid_enable_step_clipping: bool = True

hybrid_max_warmup_step_size: float = 0.1

enable_multi_start: bool = False

multi_start_n_starts: int = 10

multi_start_seed: int = 42

multi_start_sampling_strategy: str = 'latin_hypercube'

multi_start_n_workers: int = 0

multi_start_use_screening: bool = True

multi_start_screen_keep_fraction: float = 0.5

multi_start_refine_top_k: int = 3

multi_start_refinement_ftol: float = 1e-12

multi_start_degeneracy_threshold: float = 0.1

per_angle_mode: str = 'auto'

fourier_order: int = 2

fourier_auto_threshold: int = 6

constant_scaling_threshold: int = 3

enable_hierarchical: bool = True

hierarchical_max_outer_iterations: int = 5

hierarchical_outer_tolerance: float = 1e-06

hierarchical_physical_max_iterations: int = 100

hierarchical_per_angle_max_iterations: int = 50

regularization_mode: str = 'relative'

group_variance_lambda: float = 1.0

regularization_target_cv: float = 0.1

regularization_target_contribution: float = 0.1

regularization_max_cv: float = 0.2

regularization_auto_tune_lambda: bool = True

enable_gradient_monitoring: bool = True

gradient_ratio_threshold: float = 0.01

gradient_consecutive_triggers: int = 5

gradient_collapse_response: str = 'hierarchical'

enable_cmaes: bool = False

cmaes_preset: str = 'cmaes'

cmaes_max_generations: int | None = None

cmaes_popsize: int | None = None

cmaes_sigma: float = 0.5

cmaes_sigma_warmstart: float = 0.05

cmaes_warmstart_auto_skip: bool = True

cmaes_warmstart_skip_threshold: float = 5.0

cmaes_tol_fun: float = 1e-08

cmaes_tol_x: float = 1e-08

cmaes_restart_strategy: str = 'bipop'

cmaes_max_restarts: int = 9

cmaes_population_batch_size: int | None = None

cmaes_data_chunk_size: int | None = None

cmaes_refine_with_nlsq: bool = True

cmaes_auto_select: bool = True

cmaes_scale_threshold: float = 1000.0

cmaes_memory_limit_gb: float = 8.0

cmaes_refinement_workflow: str = 'auto'

cmaes_refinement_ftol: float = 1e-10

cmaes_refinement_xtol: float = 1e-10

cmaes_refinement_gtol: float = 1e-10

cmaes_refinement_max_nfev: int = 500

cmaes_refinement_loss: str = 'linear'

cmaes_normalize: bool = True

cmaes_normalization_epsilon: float = 1e-12

enable_quality_validation: bool = True

quality_reduced_chi_squared_threshold: float = 10.0

quality_warn_on_max_restarts: bool = True

quality_warn_on_bounds_hit: bool = True

quality_warn_on_convergence_failure: bool = True

quality_bounds_tolerance: float = 1e-09

classmethod from_dict(config_dict)

Create NLSQConfig from configuration dictionary.

Parameters:

config_dict (dict[str, Any]) – NLSQ configuration dictionary from ConfigManager.

Returns:

Validated configuration object.

Return type:

NLSQConfig

classmethod from_yaml(yaml_path)

Create NLSQConfig from YAML configuration file (T099).

This is the recommended single entry point for loading NLSQ configuration. It reads the YAML file, extracts the optimization.nlsq section, and creates a validated NLSQConfig object.

Parameters:

yaml_path (str) – Path to YAML configuration file.

Returns:

Validated configuration object.

Return type:

NLSQConfig

Raises:

• FileNotFoundError – If the YAML file does not exist.

• ValueError – If the YAML file is invalid or missing required sections.

Examples

>>> config = NLSQConfig.from_yaml("homodyne_config.yaml")
>>> print(config.loss)
soft_l1

validate()

Validate configuration values.

Returns:

List of validation error messages (empty if valid).

Return type:

list[str]

is_valid()

Check if configuration is valid.

Returns:

True if configuration has no validation errors.

Return type:

bool

to_dict()

Convert configuration to dictionary.

Returns:

Configuration as dictionary.

Return type:

dict[str, Any]

to_workflow_kwargs()

Convert settings to kwargs for NLSQ’s curve_fit().

Maps NLSQConfig settings to NLSQ 0.6.4+ curve_fit() parameters. Note: Homodyne uses curve_fit() directly, not the fit() unified API.

Returns:

Kwargs for curve_fit() (ftol, gtol, xtol, max_nfev, loss).

Return type:

dict[str, Any]

Notes

NLSQ 0.6.3+ Changes: - Simplified to 3 workflows: “auto”, “auto_global”, “hpc” - Old presets (“streaming”, “standard”) were removed - WorkflowSelector was removed; use MemoryBudgetSelector instead - Homodyne uses its own select_nlsq_strategy() for memory selection

The ‘goal’ parameter can be passed to NLSQ’s fit() API but homodyne uses curve_fit() directly, so goal is handled internally.

Example

>>> config = NLSQConfig.from_dict(yaml_config)
>>> kwargs = config.to_workflow_kwargs()
>>> result = fitter.curve_fit(f, xdata, ydata, **kwargs)

__init__(workflow='auto', goal='quality', loss='soft_l1', trust_region_scale=1.0, max_iterations=1000, ftol=1e-08, xtol=1e-08, gtol=1e-08, x_scale='jac', x_scale_map=None, enable_diagnostics=True, enable_streaming=True, streaming_chunk_size=50000, enable_stratified=True, target_chunk_size=100000, enable_recovery=True, max_recovery_attempts=3, enable_progress_bar=True, verbose=1, log_iteration_interval=10, enable_hybrid_streaming=True, hybrid_normalize=True, hybrid_normalization_strategy='auto', hybrid_warmup_iterations=200, hybrid_max_warmup_iterations=500, hybrid_warmup_learning_rate=0.001, hybrid_gauss_newton_max_iterations=100, hybrid_gauss_newton_tol=1e-08, hybrid_chunk_size=10000, hybrid_trust_region_initial=1.0, hybrid_regularization_factor=1e-10, hybrid_enable_checkpoints=True, hybrid_checkpoint_frequency=100, hybrid_validate_numerics=True, hybrid_enable_warm_start_detection=True, hybrid_warm_start_threshold=0.01, hybrid_enable_adaptive_warmup_lr=True, hybrid_warmup_lr_refinement=1e-06, hybrid_warmup_lr_careful=1e-05, hybrid_enable_cost_guard=True, hybrid_cost_increase_tolerance=0.05, hybrid_enable_step_clipping=True, hybrid_max_warmup_step_size=0.1, enable_multi_start=False, multi_start_n_starts=10, multi_start_seed=42, multi_start_sampling_strategy='latin_hypercube', multi_start_n_workers=0, multi_start_use_screening=True, multi_start_screen_keep_fraction=0.5, multi_start_refine_top_k=3, multi_start_refinement_ftol=1e-12, multi_start_degeneracy_threshold=0.1, per_angle_mode='auto', fourier_order=2, fourier_auto_threshold=6, constant_scaling_threshold=3, enable_hierarchical=True, hierarchical_max_outer_iterations=5, hierarchical_outer_tolerance=1e-06, hierarchical_physical_max_iterations=100, hierarchical_per_angle_max_iterations=50, regularization_mode='relative', group_variance_lambda=1.0, regularization_target_cv=0.1, regularization_target_contribution=0.1, regularization_max_cv=0.2, regularization_auto_tune_lambda=True, enable_gradient_monitoring=True, gradient_ratio_threshold=0.01, gradient_consecutive_triggers=5, gradient_collapse_response='hierarchical', enable_cmaes=False, cmaes_preset='cmaes', cmaes_max_generations=None, cmaes_popsize=None, cmaes_sigma=0.5, cmaes_sigma_warmstart=0.05, cmaes_warmstart_auto_skip=True, cmaes_warmstart_skip_threshold=5.0, cmaes_tol_fun=1e-08, cmaes_tol_x=1e-08, cmaes_restart_strategy='bipop', cmaes_max_restarts=9, cmaes_population_batch_size=None, cmaes_data_chunk_size=None, cmaes_refine_with_nlsq=True, cmaes_auto_select=True, cmaes_scale_threshold=1000.0, cmaes_memory_limit_gb=8.0, cmaes_refinement_workflow='auto', cmaes_refinement_ftol=1e-10, cmaes_refinement_xtol=1e-10, cmaes_refinement_gtol=1e-10, cmaes_refinement_max_nfev=500, cmaes_refinement_loss='linear', cmaes_normalize=True, cmaes_normalization_epsilon=1e-12, enable_quality_validation=True, quality_reduced_chi_squared_threshold=10.0, quality_warn_on_max_restarts=True, quality_warn_on_bounds_hit=True, quality_warn_on_convergence_failure=True, quality_bounds_tolerance=1e-09, _validation_errors=<factory>)

Key Classes

homodyne.optimization.nlsq.config.NLSQConfig	Configuration for NLSQ (Nonlinear Least Squares) optimization.
homodyne.optimization.nlsq.config.safe_float	Convert value to float safely, returning default on failure.
homodyne.optimization.nlsq.config.safe_int	Convert value to int safely, returning default on failure.

Configuration Entry Point

Use NLSQConfig.from_yaml() as the single entry point for loading NLSQ configuration:

from homodyne.optimization.nlsq.config import NLSQConfig

# Load configuration from YAML file
config = NLSQConfig.from_yaml("config.yaml")

# Access configuration values
print(f"Tolerance: {config.tolerance}")
print(f"Max iterations: {config.max_iterations}")

Configuration Utilities (Deprecated)

Deprecated since version 2.14.0: Use homodyne.optimization.nlsq.config instead. The safe_float, safe_int, and safe_bool utilities have been moved to config.py.

Note

config_utils was merged into homodyne.optimization.nlsq.config. See safe_float() and safe_int().

NLSQAdapterBase

Abstract base class providing shared functionality for NLSQAdapter and NLSQWrapper. This enables code reuse and consistent interfaces across both adapter implementations.

Abstract base class for NLSQ adapters (FR-012).

Provides shared methods for NLSQAdapter and NLSQWrapper to reduce code duplication.

Created as part of architecture refactoring (T059-T061).

class homodyne.optimization.nlsq.adapter_base.NLSQAdapterBase

Bases: ABC

Abstract base class for NLSQ optimization adapters.

Provides shared methods for data preparation, validation, result building, error handling, bounds setup, and covariance computation.

Subclasses must implement the fit() method.

abstractmethod fit(*args, **kwargs)

Fit the model to data.

Must be implemented by subclasses.

Return type:

Any

Key Classes

homodyne.optimization.nlsq.adapter_base.NLSQAdapterBase

Abstract base class for NLSQ optimization adapters.

Shared Methods

The NLSQAdapterBase provides these common methods:

• _prepare_data(): Flatten and validate input data

• _validate_input(): Input validation with shape and type checking

• _build_result(): Construct optimization result objects

• _handle_error(): Error handling with recovery actions

• _setup_bounds(): Bounds configuration and validation

• _compute_covariance(): Covariance matrix computation from Jacobian

Input and Result Validation

Validation utilities extracted from wrapper.py for independent testing and reuse.

Input Validator

Input validation for NLSQ optimization (T079).

Extracted from wrapper.py as part of architecture refactoring. Enhanced with structured logging for T027.

class homodyne.optimization.nlsq.validation.input_validator.InputValidator

Bases: object

Validator for NLSQ optimization input data.

Validates input arrays, bounds, initial parameters, and configuration before optimization begins.

__init__(strict_mode=True)

Initialize InputValidator.

Parameters:

strict_mode (bool) – If True, raise errors on validation failures. If False, log warnings but continue.

validate_all(xdata, ydata, initial_params, bounds)

Validate all input data.

Parameters:

• xdata (ndarray) – Independent variable data (t1, t2, phi)

• ydata (ndarray) – Dependent variable data (g2 values)

• initial_params (ndarray) – Initial parameter guess

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds (lower, upper)

Returns:

True if all validation passes, False otherwise

Return type:

bool

property validation_errors: list[str]

Get list of validation errors from last validate_all() call.

homodyne.optimization.nlsq.validation.input_validator.validate_array_dimensions(xdata, ydata)

Validate that xdata and ydata have compatible dimensions.

Parameters:

• xdata (ndarray) – Independent variable data

• ydata (ndarray) – Dependent variable data

Returns:

True if dimensions are compatible

Return type:

bool

homodyne.optimization.nlsq.validation.input_validator.validate_no_nan_inf(arr, name, iteration=None, context=None)

Validate that array contains no NaN or Inf values (T027).

Parameters:

• arr (ndarray) – Array to validate

• name (str) – Name for logging

• iteration (int | None) – Current iteration number for context

• context (dict[str, Any] | None) – Additional context for logging

Returns:

True if array contains only finite values

Return type:

bool

homodyne.optimization.nlsq.validation.input_validator.validate_bounds_consistency(bounds, initial_params)

Validate that bounds are consistent.

Parameters:

• bounds (tuple[ndarray, ndarray]) – (lower, upper) bounds arrays

• initial_params (ndarray) – Initial parameter values

Returns:

True if bounds are consistent

Return type:

bool

homodyne.optimization.nlsq.validation.input_validator.validate_initial_params(initial_params, bounds)

Validate that initial parameters are within bounds.

Parameters:

• initial_params (ndarray) – Initial parameter values

• bounds (tuple[ndarray, ndarray] | None) – (lower, upper) bounds arrays, or None for unbounded

Returns:

True if params are within bounds

Return type:

bool

Key Functions

homodyne.optimization.nlsq.validation.input_validator.InputValidator	Validator for NLSQ optimization input data.
homodyne.optimization.nlsq.validation.input_validator.validate_array_dimensions	Validate that xdata and ydata have compatible dimensions.
homodyne.optimization.nlsq.validation.input_validator.validate_no_nan_inf	Validate that array contains no NaN or Inf values (T027).
homodyne.optimization.nlsq.validation.input_validator.validate_bounds_consistency	Validate that bounds are consistent.
homodyne.optimization.nlsq.validation.input_validator.validate_initial_params	Validate that initial parameters are within bounds.

Result Validator

Result validation for NLSQ optimization (T080).

Extracted from wrapper.py as part of architecture refactoring.

class homodyne.optimization.nlsq.validation.result_validator.ResultValidator

Bases: object

Validator for NLSQ optimization results.

Validates optimized parameters, covariance matrices, and result consistency.

__init__(strict_mode=False)

Initialize ResultValidator.

Parameters:

strict_mode (bool) – If True, raise errors on validation failures. If False, log warnings but continue.

validate_all(params, covariance, bounds, chi_squared=None)

Validate all result components.

Parameters:

• params (ndarray) – Optimized parameter values

• covariance (ndarray | None) – Parameter covariance matrix

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds (lower, upper)

• chi_squared (float | None) – Chi-squared value for quality check

Returns:

True if all validation passes, False otherwise

Return type:

bool

property validation_warnings: list[str]

Get list of validation warnings from last validate_all() call.

homodyne.optimization.nlsq.validation.result_validator.validate_optimized_params(params, bounds, tolerance=1e-10)

Validate that optimized parameters are finite and within bounds.

Parameters:

• params (ndarray) – Optimized parameter values

• bounds (tuple[ndarray, ndarray] | None) – (lower, upper) bounds arrays

• tolerance (float) – Tolerance for boundary violations

Returns:

True if params are valid

Return type:

bool

homodyne.optimization.nlsq.validation.result_validator.validate_covariance(covariance, n_params)

Validate covariance matrix properties.

Parameters:

• covariance (ndarray) – Parameter covariance matrix

• n_params (int) – Expected number of parameters

Returns:

True if covariance is valid

Return type:

bool

homodyne.optimization.nlsq.validation.result_validator.validate_result_consistency(params, chi_squared)

Validate consistency of optimization result.

Parameters:

• params (ndarray) – Optimized parameter values

• chi_squared (float) – Chi-squared value

Returns:

True if result is consistent

Return type:

bool

Key Functions

homodyne.optimization.nlsq.validation.result_validator.ResultValidator	Validator for NLSQ optimization results.
homodyne.optimization.nlsq.validation.result_validator.validate_optimized_params	Validate that optimized parameters are finite and within bounds.
homodyne.optimization.nlsq.validation.result_validator.validate_covariance	Validate covariance matrix properties.
homodyne.optimization.nlsq.validation.result_validator.validate_result_consistency	Validate consistency of optimization result.

Fit Quality Validator

Post-optimization quality checks with configurable thresholds.

Fit quality validation for NLSQ results (T056).

Provides post-optimization quality checks with configurable thresholds. Logs warnings for potential issues but does not raise exceptions.

Usage:

>>> from homodyne.optimization.nlsq.validation.fit_quality import (
...     FitQualityConfig,
...     validate_fit_quality,
... )
>>> config = FitQualityConfig(reduced_chi_squared_threshold=10.0)
>>> report = validate_fit_quality(result, bounds=bounds, config=config)
>>> if not report.passed:
...     print(f"Warnings: {report.warnings}")

class homodyne.optimization.nlsq.validation.fit_quality.FitQualityConfig

Bases: object

Configuration for fit quality validation.

enable

Whether to enable quality validation. Default: True.

Type:

bool

reduced_chi_squared_threshold

Warn if reduced chi-squared exceeds this. Default: 10.0.

Type:

float

chi2_good_threshold

Reduced chi-squared below which fit is classified as “good”. Default: 2.0.

Type:

float

chi2_acceptable_threshold

Reduced chi-squared below which fit is classified as “acceptable”. Default: 5.0.

Type:

float

min_parameter_significance

Minimum parameter/uncertainty ratio for significance. Default: 2.0.

Type:

float

max_condition_number

Maximum covariance matrix condition number. Default: 1e12.

Type:

float

warn_on_max_restarts

Warn if CMA-ES reached max_restarts. Default: True.

Type:

bool

warn_on_bounds_hit

Warn if physical parameters hit bounds. Default: True.

Type:

bool

warn_on_convergence_failure

Warn if convergence_status indicates failure. Default: True.

Type:

bool

bounds_tolerance

Tolerance for “at bounds” detection. Default: 1e-9.

Type:

float

enable: bool = True

reduced_chi_squared_threshold: float = 10.0

chi2_good_threshold: float = 2.0

chi2_acceptable_threshold: float = 5.0

min_parameter_significance: float = 2.0

max_condition_number: float = 1000000000000.0

warn_on_max_restarts: bool = True

warn_on_bounds_hit: bool = True

warn_on_convergence_failure: bool = True

bounds_tolerance: float = 1e-09

classmethod from_validation_config(validation_config)

Create FitQualityConfig from an NLSQValidationConfig dict.

Parameters:

validation_config (dict[str, Any] | None) – Dictionary with keys from NLSQValidationConfig TypedDict. If None, returns defaults.

Returns:

Configuration with values from the dict, falling back to defaults.

Return type:

FitQualityConfig

__init__(enable=True, reduced_chi_squared_threshold=10.0, chi2_good_threshold=2.0, chi2_acceptable_threshold=5.0, min_parameter_significance=2.0, max_condition_number=1000000000000.0, warn_on_max_restarts=True, warn_on_bounds_hit=True, warn_on_convergence_failure=True, bounds_tolerance=1e-09)

class homodyne.optimization.nlsq.validation.fit_quality.FitQualityReport

Bases: object

Report from fit quality validation.

passed

True if no warnings were generated.

Type:

bool

warnings

List of warning messages.

Type:

list[str]

checks_performed

Which checks were performed and their pass/fail status.

Type:

dict[str, bool]

passed: bool = True

warnings: list[str]

checks_performed: dict[str, bool]

to_dict()

Convert to dictionary for saving in results.

Return type:

dict[str, Any]

__init__(passed=True, warnings=<factory>, checks_performed=<factory>)

homodyne.optimization.nlsq.validation.fit_quality.validate_fit_quality(result, bounds=None, config=None, param_labels=None)

Validate fit quality and log warnings.

Parameters:

• result (Any) – NLSQ optimization result.

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds (lower, upper) for bounds checking.

• config (FitQualityConfig | None) – Validation configuration. Uses defaults if None.

• param_labels (list[str] | None) – Parameter labels for identifying physical vs scaling params.

Returns:

Validation report with warnings and check results.

Return type:

FitQualityReport

Key Classes

homodyne.optimization.nlsq.validation.fit_quality.FitQualityConfig	Configuration for fit quality validation.
homodyne.optimization.nlsq.validation.fit_quality.FitQualityReport	Report from fit quality validation.
homodyne.optimization.nlsq.validation.fit_quality.validate_fit_quality	Validate fit quality and log warnings.

Quality Checks

1. Reduced χ² threshold: Warns if χ²_reduced > threshold (default 10.0)

2. CMA-ES convergence: Warns if CMA-ES reached max_restarts without converging

3. Physical parameters at bounds: Warns if D₀, α, γ̇₀, etc. hit their bounds

4. Convergence status: Warns if optimization failed or hit max iterations

Usage Example

from homodyne.optimization.nlsq.validation import InputValidator, ResultValidator
import numpy as np

# Input validation
validator = InputValidator(strict_mode=True)
xdata = np.random.rand(1000, 3)
ydata = np.random.rand(1000)
initial = np.array([1000.0, 0.8, 100.0])
bounds = (np.array([100, 0, 0]), np.array([10000, 2, 1000]))

is_valid = validator.validate_all(xdata, ydata, initial, bounds)

# Result validation
result_validator = ResultValidator(strict_mode=False)
optimized = np.array([1234.5, 0.85, 150.0])
covariance = np.eye(3) * 0.01

is_valid = result_validator.validate_all(optimized, covariance, bounds)

Supporting Modules

The optimization module includes several supporting utilities:

Checkpoint Manager

Checkpoint management for streaming optimization.

This module provides checkpoint save/load functionality for fault-tolerant streaming optimization. Checkpoints are stored in HDF5 format with compression and checksum validation.

Key Features: - HDF5-based checkpoint storage with compression - Checksum validation for integrity - Automatic cleanup of old checkpoints - Version compatibility checking - Fast save time (< 2 seconds target)

The CheckpointManager complements NLSQ’s built-in checkpointing by storing homodyne-specific state (batch statistics, recovery actions, best parameters).

class homodyne.optimization.checkpoint_manager.CheckpointManager

Bases: object

Manage checkpoint save/load for streaming optimization.

This class provides checkpoint management for homodyne-specific state during streaming optimization. It complements NLSQ’s built-in checkpoint functionality by storing additional metadata, batch statistics, and recovery action history.

Features: - HDF5-based checkpoint storage with compression - Checksum validation for integrity - Automatic cleanup of old checkpoints - Version compatibility checking

checkpoint_dir

Directory for checkpoint files

Type:

Path

checkpoint_frequency

Save checkpoint every N batches

Type:

int

keep_last_n

Keep last N checkpoints (default: 3)

Type:

int

enable_compression

Use HDF5 compression (default: True)

Type:

bool

Examples

>>> manager = CheckpointManager("./checkpoints", checkpoint_frequency=10)
>>> # Save checkpoint
>>> path = manager.save_checkpoint(
...     batch_idx=10,
...     parameters=params,
...     optimizer_state={'iteration': 42},
...     loss=0.123,
... )
>>> # Load checkpoint
>>> data = manager.load_checkpoint(path)
>>> params = data['parameters']
>>> batch_idx = data['batch_idx']

__init__(checkpoint_dir, checkpoint_frequency=10, keep_last_n=3, enable_compression=True)

Initialize checkpoint manager.

Parameters:

• checkpoint_dir (str | Path) – Directory for checkpoint files

• checkpoint_frequency (int) – Save checkpoint every N batches, by default 10

• keep_last_n (int) – Keep last N checkpoints, by default 3

• enable_compression (bool) – Use HDF5 compression, by default True

save_checkpoint(batch_idx, parameters, optimizer_state, loss, metadata=None)

Save checkpoint to HDF5 file.

Saves checkpoint with compression, checksum validation, and version information. Target save time is < 2 seconds for typical parameter sets.

Parameters:

• batch_idx (int) – Current batch index

• parameters (ndarray) – Current parameter values

• optimizer_state (dict) – Optimizer internal state

• loss (float) – Current loss value

• metadata (dict | None) – Additional metadata (batch statistics, recovery actions, etc.)

Returns:

Path to saved checkpoint file

Return type:

Path

Raises:

NLSQCheckpointError – If checkpoint save fails

Notes

Checkpoint file naming: homodyne_state_batch_{batch_idx:04d}.h5

load_checkpoint(checkpoint_path)

Load and validate checkpoint.

Loads checkpoint from HDF5 file and validates checksum integrity.

Security: Uses pickle.loads() for optimizer state deserialization. This is safe because checkpoint files are created exclusively by save_checkpoint() with checksum validation, stored in application-controlled output directories, and the serialized bytes are embedded within HDF5 containers created by this class.

Parameters:

checkpoint_path (Path) – Path to checkpoint file

Returns:

Checkpoint data with keys: - batch_idx: int - Batch index when checkpoint was saved - parameters: np.ndarray - Parameter values - optimizer_state: dict - Optimizer internal state - loss: float - Loss value at checkpoint - metadata: dict - Additional metadata (if available) - version: str - Homodyne version - timestamp: float - Unix timestamp

Return type:

dict

Raises:

NLSQCheckpointError – If checkpoint is corrupted, invalid, or missing

find_latest_checkpoint()

Find most recent valid checkpoint.

Searches checkpoint directory for valid checkpoint files and returns the one with the highest batch index.

Returns:

Path to latest checkpoint, or None if none exist

Return type:

Path | None

Notes

Only returns checkpoints that pass validation.

cleanup_old_checkpoints()

Remove old checkpoints, keeping last N.

Keeps the most recent N checkpoints based on batch index and removes older ones to manage disk space.

Returns:

Paths of deleted checkpoints

Return type:

list[Path]

Notes

Only deletes checkpoints, never removes the keep_last_n most recent ones.

validate_checkpoint(checkpoint_path)

Validate checkpoint integrity.

Checks that checkpoint file exists, can be opened, has required fields, and passes checksum validation.

Parameters:

checkpoint_path (Path) – Path to checkpoint file

Returns:

True if valid, False otherwise

Return type:

bool

Gradient Diagnostics

Gradient Diagnostics for Parameter Scaling Optimization

This module provides tools to diagnose gradient imbalance issues and compute optimal parameter scaling factors (x_scale) for NLSQ optimization.

The Problem:

Shear parameters (gamma_dot_t0, beta, gamma_dot_t_offset) can have gradients 100-10,000× larger than diffusion parameters (D0, alpha, D_offset), causing:

• Premature convergence

• Missing fine-scale features (oscillations)

• Poor fit quality despite low chi-squared

The Solution:

Compute parameter-specific x_scale values inversely proportional to gradient magnitudes to normalize optimization steps across all parameters.

Usage:

from homodyne.optimization.gradient_diagnostics import compute_optimal_x_scale

# Compute from fitted parameters
x_scale_map = compute_optimal_x_scale(
    parameters=result.parameters,
    data=data,
    config=config,
    analysis_mode="laminar_flow"
)

# Add to config for next optimization
config.config["optimization"]["nlsq"]["x_scale_map"] = x_scale_map

Author: Homodyne Development Team Date: 2025-11-13 Version: 1.0.0

homodyne.optimization.gradient_diagnostics.compute_gradient_norms(parameters, data, config, analysis_mode)

Compute gradient L2 norms for each parameter at given point.

Parameters:

• parameters (dict[str, float]) – Dictionary of parameter values

• data (Any) – Data object with experimental data

• config (Any) – Configuration object

• analysis_mode (str) – “static_isotropic” or “laminar_flow”

Return type:

dict[str, float]

Returns:

Dictionary mapping parameter names to gradient norms

Example

>>> gradient_norms = compute_gradient_norms(
...     parameters=result.parameters,
...     data=data,
...     config=config,
...     analysis_mode="laminar_flow"
... )
>>> # Output: {'D0': 26.98, 'alpha': 42365.33, ..., 'gamma_dot_t_offset': 346934800.0}

homodyne.optimization.gradient_diagnostics.compute_optimal_x_scale(parameters, data, config, analysis_mode, baseline_params=None, safety_factor=1.0, min_scale=1e-8, max_scale=1e2)

Compute optimal x_scale map based on gradient norms.

The x_scale values are inversely proportional to gradient magnitudes, normalized so that baseline parameters have x_scale=1.0.

Parameters:

• parameters (dict[str, float]) – Dictionary of parameter values

• data (Any) – Data object with experimental data

• config (Any) – Configuration object

• analysis_mode (str) – “static_isotropic” or “laminar_flow”

• baseline_params (list[str] | None) – Parameters to use as baseline (x_scale=1.0). Default: [“D0”, “D_offset”, “phi0”]

• safety_factor (float) – Multiplicative safety factor (default: 1.0) Increase to make optimization more conservative

• min_scale (float) – Minimum allowed x_scale value (prevents division by zero)

• max_scale (float) – Maximum allowed x_scale value (prevents extreme values)

Return type:

dict[str, float]

Returns:

Dictionary mapping parameter names to x_scale values

Example

>>> x_scale_map = compute_optimal_x_scale(
...     parameters={'D0': 400.0, 'alpha': -0.014, ..., 'gamma_dot_t_offset': 0.0},
...     data=data,
...     config=config,
...     analysis_mode="laminar_flow"
... )
>>> # Output: {'D0': 1.0, 'alpha': 0.001, ..., 'gamma_dot_t_offset': 1e-7}

homodyne.optimization.gradient_diagnostics.diagnose_gradient_imbalance(parameters, data, config, analysis_mode, threshold=10.0)

Diagnose gradient imbalance and provide recommendations.

Parameters:

• parameters (dict[str, float]) – Dictionary of parameter values

• data (Any) – Data object with experimental data

• config (Any) – Configuration object

• analysis_mode (str) – “static_isotropic” or “laminar_flow”

• threshold (float) – Gradient ratio threshold for warning (default: 10.0)

Returns:

• gradient_norms: Dict[str, float] - gradient norms for each parameter

• imbalance_detected: bool - whether imbalance exceeds threshold

• max_ratio: float - maximum gradient ratio

• recommendations: Dict[str, Any] - optimization recommendations

Return type:

dict[str, Any]

Example

>>> diag = diagnose_gradient_imbalance(
...     parameters=result.parameters,
...     data=data,
...     config=config,
...     analysis_mode="laminar_flow"
... )
>>> if diag["imbalance_detected"]:
...     print(f"Gradient imbalance detected: max ratio = {diag['max_ratio']:.0f}x")
...     print("Recommendations:")
...     print(diag["recommendations"]["summary"])

homodyne.optimization.gradient_diagnostics.print_gradient_report(parameters, data, config, analysis_mode)

Print comprehensive gradient diagnostic report.

Parameters:

• parameters (dict[str, float]) – Dictionary of parameter values

• data (Any) – Data object with experimental data

• config (Any) – Configuration object

• analysis_mode (str) – “static_isotropic” or “laminar_flow”

Return type:

None

Example

>>> # After NLSQ optimization
>>> print_gradient_report(
...     parameters=result.parameters,
...     data=data,
...     config=config,
...     analysis_mode="laminar_flow"
... )
# Prints detailed gradient analysis and recommendations

Exceptions

Custom exceptions for NLSQ optimization.

This module defines a comprehensive exception hierarchy for handling errors specific to NLSQ optimization, including convergence failures, numerical instabilities, and checkpoint-related issues.

The exception hierarchy enables fine-grained error handling and recovery strategies tailored to specific failure modes.

Exception Hierarchy:

NLSQOptimizationError (base) ├── NLSQConvergenceError (convergence failures) ├── NLSQNumericalError (NaN/Inf issues) └── NLSQCheckpointError (checkpoint save/load failures)

Examples

Catching specific errors for targeted recovery:

>>> try:
...     result = optimizer.fit(data, model, p0)
... except NLSQNumericalError as e:
...     # Handle NaN/Inf with learning rate reduction
...     result = optimizer.fit(data, model, p0, learning_rate=0.5*lr)
... except NLSQConvergenceError as e:
...     # Handle convergence failure with perturbation
...     p0_perturbed = p0 * (1 + 0.01 * np.random.randn(*p0.shape))
...     result = optimizer.fit(data, model, p0_perturbed)

Using base exception for generic handling:

>>> try:
...     result = optimizer.fit(data, model, p0)
... except NLSQOptimizationError as e:
...     logger.error(f"Optimization failed: {e}")
...     # Fallback to simpler strategy
...     result = use_fallback_strategy()

Notes

All exceptions inherit from NLSQOptimizationError, enabling catch-all error handling while also supporting fine-grained recovery strategies.

The exception messages are designed to be actionable, providing specific guidance on how to address each type of failure.

See also

NLSQWrapper

Main optimization wrapper using these exceptions

homodyne.optimization.strategy

Strategy selection and fallback logic

exception homodyne.optimization.exceptions.NLSQOptimizationError

Bases: Exception

Base exception for all NLSQ optimization errors.

This is the base class for all NLSQ-related exceptions. Catching this exception will catch all optimization failures regardless of their specific cause.

message

Detailed error message

Type:

str

error_context

Additional context about the error (parameters, data characteristics, etc.)

Type:

dict

Examples

>>> try:
...     result = optimizer.fit(data, model, p0)
... except NLSQOptimizationError as e:
...     print(f"Optimization failed: {e}")
...     print(f"Context: {e.error_context}")

__init__(message, error_context=None)

Initialize base optimization error.

Parameters:

• message (str) – Detailed error message

• error_context (dict | None) – Additional context about the error

__str__()

Return formatted error message with context.

Return type:

str

exception homodyne.optimization.exceptions.NLSQConvergenceError

Bases: NLSQOptimizationError

Raised when NLSQ optimization fails to converge.

This exception indicates that the optimizer could not find a satisfactory solution within the specified constraints (maximum iterations, tolerance, etc.).

Common Causes

• Poor initial guess (p0 too far from optimum)

• Overly restrictive parameter bounds

• Insufficient maximum iterations

• Model function incompatible with data

• Local minimum trap

Recovery Strategies

1. Perturb initial guess: p0 * (1 + 0.05 * np.random.randn(*p0.shape))

2. Relax bounds: Increase parameter search space

3. Increase max iterations: Allow more optimization steps

4. Try different optimization method: Switch between ‘trf’ and ‘lm’

5. Simplify model: Use fewer parameters

iteration_count

Number of iterations completed before failure

Type:

int

final_loss

Final loss value at termination

Type:

float

parameters

Parameter values at termination

Type:

np.ndarray

Examples

>>> try:
...     result = optimizer.fit(data, model, p0, max_iter=100)
... except NLSQConvergenceError as e:
...     print(f"Failed after {e.iteration_count} iterations")
...     print(f"Final loss: {e.final_loss}")
...     # Retry with more iterations
...     result = optimizer.fit(data, model, p0, max_iter=500)

__init__(message, iteration_count=None, final_loss=None, parameters=None, error_context=None)

Initialize convergence error.

Parameters:

• message (str) – Detailed error message

• iteration_count (int | None) – Number of iterations completed

• final_loss (float | None) – Final loss value

• parameters (ndarray | None) – Parameter values at termination

• error_context (dict | None) – Additional context

exception homodyne.optimization.exceptions.NLSQNumericalError

Bases: NLSQOptimizationError

Raised for NaN/Inf numerical stability issues.

This exception indicates that the optimization encountered numerical instabilities such as NaN (Not a Number) or Inf (Infinity) values during computation.

Common Causes

• Gradient overflow/underflow

• Division by zero in model function

• Exponential overflow in parameters

• Ill-conditioned Jacobian matrix

• Learning rate too large

Detection Points

1. After gradient computation: jnp.isfinite(gradients).all()

2. After parameter update: jnp.isfinite(new_params).all()

3. After loss calculation: jnp.isfinite(loss_value)

Recovery Strategies

1. Reduce learning rate: lr = 0.5 * lr

2. Scale data: Normalize inputs to [0, 1] range

3. Add numerical stability: Use log-transform for exponentials

4. Check model function: Ensure JAX-compatible operations

5. Adjust parameter bounds: Prevent extreme values

detection_point

Where NaN/Inf was detected (‘gradient’, ‘parameter’, ‘loss’)

Type:

str

invalid_values

Description of invalid values found

Type:

list

Examples

>>> try:
...     result = optimizer.fit(data, model, p0)
... except NLSQNumericalError as e:
...     if e.detection_point == 'gradient':
...         # Reduce learning rate
...         result = optimizer.fit(data, model, p0, learning_rate=0.01)
...     elif e.detection_point == 'parameter':
...         # Tighten bounds
...         bounds = (lower * 0.8, upper * 0.8)
...         result = optimizer.fit(data, model, p0, bounds=bounds)

__init__(message, detection_point=None, invalid_values=None, error_context=None)

Initialize numerical error.

Parameters:

• message (str) – Detailed error message

• detection_point (str | None) – Where NaN/Inf was detected

• invalid_values (list | None) – Description of invalid values

• error_context (dict | None) – Additional context

exception homodyne.optimization.exceptions.NLSQCheckpointError

Bases: NLSQOptimizationError

Raised for checkpoint save/load/resume failures.

This exception indicates that the streaming optimizer encountered an error while saving checkpoints, loading checkpoints, or resuming from a checkpoint.

Common Causes

• Checkpoint file corrupted

• Insufficient disk space

• Invalid checkpoint path

• HDF5 file lock conflict

• Version mismatch in checkpoint format

• Missing checkpoint metadata

Recovery Strategies

1. Disable checkpoints: config.enable_checkpoints = False

2. Change checkpoint directory: Use different storage location

3. Clear old checkpoints: Remove corrupted checkpoint files

4. Start fresh: config.resume_from_checkpoint = False

5. Reduce checkpoint frequency: Save less often to avoid I/O issues

checkpoint_path

Path to the checkpoint file involved

Type:

str

operation

Operation that failed (‘save’, ‘load’, ‘resume’, ‘validate’)

Type:

str

io_error

Original I/O exception if available

Type:

Exception

Examples

>>> try:
...     config = HybridStreamingConfig(enable_checkpoints=True)
...     optimizer = AdaptiveHybridStreamingOptimizer(config)
...     result = optimizer.fit(data, model, p0)
... except NLSQCheckpointError as e:
...     if e.operation == 'load':
...         # Start fresh if checkpoint is corrupted
...         config = HybridStreamingConfig(enable_checkpoints=False)
...         optimizer = AdaptiveHybridStreamingOptimizer(config)
...         result = optimizer.fit(data, model, p0)
...     elif e.operation == 'save':
...         # Continue without checkpoints
...         config = HybridStreamingConfig(enable_checkpoints=False)
...         optimizer = AdaptiveHybridStreamingOptimizer(config)
...         result = optimizer.fit(data, model, p0)

__init__(message, checkpoint_path=None, operation=None, io_error=None, error_context=None)

Initialize checkpoint error.

Parameters:

• message (str) – Detailed error message

• checkpoint_path (str | None) – Path to checkpoint file

• operation (str | None) – Operation that failed

• io_error (Exception | None) – Original I/O exception

• error_context (dict | None) – Additional context

Recovery Strategies

Error recovery strategies for NLSQ optimization failures.

This module defines error-specific recovery strategies that can be applied when optimization encounters failures. Each error type has a prioritized list of recovery actions to attempt.

class homodyne.optimization.recovery_strategies.RecoveryStrategyApplicator

Bases: object

Apply recovery strategies for optimization failures.

This class implements various recovery strategies that can be applied when optimization fails. Strategies are error-type specific and are applied in a prioritized order.

Parameters:

max_retries (int) – Maximum number of retry attempts per batch, by default 2

Examples

>>> applicator = RecoveryStrategyApplicator(max_retries=2)
>>> error = NLSQConvergenceError("Failed to converge")
>>> strategy_name, modified_params = applicator.get_recovery_strategy(
...     error, params, attempt=0
... )
>>> # strategy_name is "perturb_parameters"
>>> # modified_params has 5% random noise added

__init__(max_retries=2, seed=42)

Initialize recovery strategy applicator.

Parameters:

• max_retries (int) – Maximum retry attempts, by default 2

• seed (int) – RNG seed for reproducible perturbations, by default 42

get_recovery_strategy(error, params, attempt, bounds=None)

Get recovery strategy for the given error and attempt.

Parameters:

• error (Exception) – The exception that was raised

• params (ndarray) – Current parameter values

• attempt (int) – Retry attempt number (0-indexed)

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds (lower, upper), by default None

Returns:

(strategy_name, modified_params) if strategy available, else None

Return type:

tuple[str, ndarray] | None

should_retry(attempt)

Check if another retry attempt should be made.

Parameters:

attempt (int) – Current attempt number (0-indexed)

Returns:

True if should retry, False if max retries exhausted

Return type:

bool

Batch Statistics

Batch-level statistics tracking for streaming optimization.

Batch-level statistics tracking for streaming optimization.

This module provides a circular buffer for tracking batch-level optimization statistics, success rates, and error distributions during streaming optimization.

class homodyne.optimization.batch_statistics.BatchStatistics

Bases: object

Circular buffer for tracking batch-level statistics.

Maintains statistics for the most recent N batches (default 100) to provide running averages and trends without unbounded memory growth.

buffer

Circular buffer storing batch records (max_size most recent)

Type:

deque

total_batches

Total number of batches processed (all time)

Type:

int

total_successes

Total number of successful batches (all time)

Type:

int

total_failures

Total number of failed batches (all time)

Type:

int

error_counts

Count of each error type encountered (all time)

Type:

dict

Examples

>>> stats = BatchStatistics(max_size=100)
>>> stats.record_batch(
...     batch_idx=0,
...     success=True,
...     loss=0.123,
...     iterations=50,
...     recovery_actions=[]
... )
>>> stats.get_success_rate()
1.0

__init__(max_size=100)

Initialize batch statistics tracker.

Parameters:

max_size (int) – Maximum number of batches to keep in circular buffer, by default 100

record_batch(batch_idx, success, loss, iterations, recovery_actions, error_type=None)

Record statistics for a single batch.

Parameters:

• batch_idx (int) – Batch index (0-indexed)

• success (bool) – Whether batch optimization succeeded

• loss (float) – Final loss value for this batch

• iterations (int) – Number of iterations performed

• recovery_actions (list[str]) – List of recovery actions applied (if any)

• error_type (str | None) – Type of error encountered (if failed), by default None

Return type:

None

get_success_rate()

Calculate success rate from recent batches in buffer.

Returns:

Success rate (0.0 to 1.0) from recent batches. Returns 1.0 when no batches have been recorded yet (optimistic prior) so that quality gates do not falsely reject the first batch. Callers that need to distinguish “no data yet” should check BatchStatistics.total_batches.

Return type:

float

get_average_loss()

Calculate average loss from recent successful batches.

Returns:

Average loss from successful batches in buffer

Return type:

float

get_average_iterations()

Calculate average iterations from recent batches.

Returns:

Average number of iterations per batch

Return type:

float

get_statistics()

Return comprehensive statistics dictionary.

Returns:

Dictionary containing: - total_batches: Total batches processed (all time) - total_successes: Total successful batches (all time) - total_failures: Total failed batches (all time) - success_rate: Success rate from recent batches - average_loss: Average loss from recent successful batches - average_iterations: Average iterations per batch - error_distribution: Dictionary of error type counts - recent_batches: List of recent batch records

Return type:

dict[str, Any]

__repr__()

Return string representation of statistics.

Return type:

str

Numerical Validation

Validation functions to detect numerical issues (NaN, Inf, bounds violations) at critical points during optimization.

Numerical validation for optimization at critical points.

This module provides validation functions to detect numerical issues (NaN, Inf, bounds violations) at three critical points during optimization: 1. After gradient computation 2. After parameter update 3. After loss calculation

These validations help catch numerical instabilities early and enable targeted recovery strategies.

class homodyne.optimization.numerical_validation.NumericalValidator

Bases: object

Validator for numerical stability at critical optimization points.

This class provides methods to validate numerical values at three critical points: gradients, parameters, and loss values. Detection of NaN/Inf enables targeted recovery strategies.

enable_validation

Whether to perform validation (can disable for speed)

Type:

bool

bounds

Parameter bounds (lower, upper) for bounds checking

Type:

tuple of np.ndarray or None

Examples

>>> validator = NumericalValidator(enable_validation=True)
>>> try:
...     validator.validate_gradients(gradients)
...     validator.validate_parameters(params, bounds)
...     validator.validate_loss(loss_value)
... except NLSQNumericalError as e:
...     print(f"Numerical error at {e.detection_point}")

__init__(enable_validation=True, bounds=None)

Initialize numerical validator.

Parameters:

• enable_validation (bool) – Whether to perform validation, by default True

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds (lower, upper), by default None

validate_gradients(gradients)

Validate gradients for NaN/Inf after Jacobian computation.

This is validation point 1: Gradients can become non-finite due to overflow in the model function or ill-conditioned Jacobian.

Parameters:

gradients (Any) – Gradient values to validate

Raises:

NLSQNumericalError – If gradients contain NaN or Inf values

Return type:

None

validate_parameters(parameters, bounds=None)

Validate parameters for NaN/Inf and bounds violations after update.

This is validation point 2: Parameters can become non-finite after update steps, especially with aggressive step sizes.

Parameters:

• parameters (Any) – Parameter values to validate

• bounds (tuple[ndarray, ndarray] | None) – Parameter bounds (lower, upper), overrides instance bounds

Raises:

NLSQNumericalError – If parameters contain NaN or Inf values

Return type:

None

validate_loss(loss_value)

Validate loss value for NaN/Inf after loss computation.

This is validation point 3: Loss can become non-finite due to overflow in residual computation or invalid parameter values.

Parameters:

loss_value (Any) – Loss value to validate

Raises:

NLSQNumericalError – If loss is NaN or Inf

Return type:

None

set_bounds(bounds)

Update parameter bounds for validation.

Parameters:

bounds (tuple[ndarray, ndarray]) – New parameter bounds (lower, upper)

Return type:

None

disable()

Disable validation for performance-critical sections.

Return type:

None

enable()

Re-enable validation after disabling.

Return type:

None

Usage Examples

NLSQ Optimization:

from homodyne.optimization import fit_nlsq_jax

result = fit_nlsq_jax(
    t1=t1,
    t2=t2,
    c2=c2,
    q2_phi_t1_t2=q2_phi_t1_t2,
    phi_rad=phi_rad,
    initial_params=initial_params,
    mode="static"
)

print(f"Best-fit D0: {result.parameters[0]:.2f}")

CMC Optimization:

from homodyne.optimization import fit_mcmc_jax, CMCConfig

config = CMCConfig(
    num_warmup=1000,
    num_samples=2000,
    num_chains=4
)

result = fit_mcmc_jax(
    t1=t1,
    t2=t2,
    c2=c2,
    q2_phi_t1_t2=q2_phi_t1_t2,
    phi_rad=phi_rad,
    initial_params=initial_params,
    mode="static",
    config=config
)

# Access posterior samples
print(result.summary())

See Also

• homodyne.core - Core physics and computation

• homodyne.config - Parameter management

• External: NLSQ Package Documentation

• External: NumPyro Documentation