.. _api-nlsq-wrapper:

====================================
homodyne.optimization.nlsq — Wrapper
====================================

``NLSQWrapper`` is the **stable fallback adapter** for advanced NLSQ
optimisation. It is preferred over :class:`~homodyne.optimization.nlsq.adapter.NLSQAdapter`
when:

- Datasets exceed 100 M points and require custom streaming/chunking strategies
- ``laminar_flow`` mode uses many phi angles (> 6) where full anti-degeneracy
  control is required
- Custom parameter transforms or advanced 3-attempt recovery mechanisms are needed
- Production stability is the primary concern

.. note::

   For most new code, use :func:`~homodyne.optimization.nlsq.core.fit_nlsq_jax`
   with ``use_adapter=True`` (default). ``NLSQWrapper`` is invoked automatically
   as a fallback if ``NLSQAdapter`` fails.

----

Key Differences from NLSQAdapter
---------------------------------

.. list-table::
   :widths: 30 35 35
   :header-rows: 1

   * - Capability
     - NLSQWrapper
     - NLSQAdapter
   * - Model caching
     - None (rebuild each fit)
     - WeakValueDictionary cache
   * - JIT compilation
     - Manual
     - Automatic
   * - Anti-degeneracy control
     - Full (all layers)
     - Via ``fit()`` delegation
   * - Streaming (>100 M pts)
     - Full custom chunking
     - Via NLSQ package
   * - Recovery strategy
     - 3-attempt retry
     - NLSQ native
   * - Per-angle stratification
     - Angle-stratified chunking
     - Automatic

----

Out-of-Core Routing
--------------------

When a dataset exceeds the memory threshold (default 75% of RAM), NLSQWrapper
routes to an out-of-core accumulation solver that processes data in chunks. This
path has a simplified Levenberg-Marquardt implementation that **lacks**
anti-degeneracy layers, hierarchical optimization, shear weighting, and adaptive
regularization.

.. important::

   **Effective parameter count pre-check**: Before memory routing, the wrapper
   queries the ``per_angle_mode`` and computes the *effective* parameter count.
   For ``auto`` mode with n_phi ≥ 3, this is 9 (7 physical + 2 averaged scaling)
   rather than the 53 expanded parameters. This prevents auto_averaged datasets
   from being unnecessarily routed to the out-of-core path.

   Without this pre-check, memory is overestimated by ~6× (53 vs 9 params),
   triggering out-of-core for datasets that fit comfortably in RAM.

**Convergence criteria**: The out-of-core solver uses multi-criteria convergence:

- ``xtol = 1e-6``: per-component maximum parameter change
- ``ftol = 1e-6``: relative cost function change

Both must be satisfied simultaneously. This replaces the previous norm-based
criterion (``||step||/||params|| < 1e-4``) which was scale-sensitive and produced
false single-iteration convergence when large-magnitude parameters (e.g.,
D₀ ~ 19000) dominated the norm.

----

NLSQWrapper
-----------

.. autoclass:: homodyne.optimization.nlsq.wrapper.NLSQWrapper
   :members:
   :undoc-members:
   :show-inheritance:
   :special-members: __init__

----

NLSQConfig
----------

.. _api-nlsq-config:

The ``NLSQConfig`` dataclass controls all NLSQ optimisation settings.
It is populated from the YAML configuration file via ``NLSQConfig.from_yaml()``.

.. autoclass:: homodyne.optimization.nlsq.config.NLSQConfig
   :members:
   :undoc-members:
   :show-inheritance:
   :special-members: __init__

HybridRecoveryConfig
~~~~~~~~~~~~~~~~~~~~

.. autoclass:: homodyne.optimization.nlsq.config.HybridRecoveryConfig
   :members:
   :undoc-members:
   :show-inheritance:

Safe type conversion utilities
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. autofunction:: homodyne.optimization.nlsq.config.safe_float
.. autofunction:: homodyne.optimization.nlsq.config.safe_int

----

Usage Examples
--------------

Direct usage of NLSQWrapper
~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

   from homodyne.optimization.nlsq.wrapper import NLSQWrapper
   from homodyne.core.fitting import ParameterSpace
   from homodyne.config.manager import ConfigManager

   config = ConfigManager("config.yaml").config
   param_space = ParameterSpace.from_config(config)

   wrapper = NLSQWrapper(config=config)
   result = wrapper.fit(data=data, parameter_space=param_space)

   print(result.parameters)
   print(result.residuals)

Large dataset (>100 M points) with streaming
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: yaml

   optimization:
     nlsq:
       memory_fraction: 0.75
       # Or set explicit memory threshold:
       # memory_threshold_gb: 48

.. code-block:: python

   from homodyne.optimization.nlsq import fit_nlsq_jax

   # NLSQWrapper handles streaming automatically when memory threshold is exceeded
   result = fit_nlsq_jax(data, config, use_adapter=False)

laminar_flow with many phi angles and full anti-degeneracy
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code-block:: python

   from homodyne.optimization.nlsq.wrapper import NLSQWrapper

   # NLSQWrapper applies all anti-degeneracy layers when n_phi > 6
   wrapper = NLSQWrapper(config=config)
   result = wrapper.fit(
       data=data,
       parameter_space=param_space,
       analysis_mode="laminar_flow",
       per_angle_mode="individual",   # Full per-angle scaling
   )

3-attempt error recovery
~~~~~~~~~~~~~~~~~~~~~~~~~

``NLSQWrapper`` retries automatically on solver failure, applying progressively
conservative settings per retry:

- Attempt 1: Original settings
- Attempt 2: 0.5× learning rate, 2× regularisation, 0.5× trust region
- Attempt 3: 0.25× learning rate, 4× regularisation, 0.25× trust region

.. code-block:: python

   from homodyne.optimization.nlsq.config import HybridRecoveryConfig

   recovery = HybridRecoveryConfig(
       max_retries=3,
       lr_decay=0.5,
       lambda_growth=2.0,
       trust_decay=0.5,
   )
   print(recovery.get_retry_settings(attempt=2))
   # {"lr_multiplier": 0.25, "lambda_multiplier": 4.0, "trust_multiplier": 0.25}

----

.. automodule:: homodyne.optimization.nlsq.wrapper
   :members:
   :undoc-members:
   :show-inheritance:
   :exclude-members: NLSQWrapper

----

Extracted Strategy Modules
---------------------------

The following modules were extracted from ``wrapper.py`` to improve
maintainability. ``NLSQWrapper`` delegates to these internally; they are not
part of the public API.

Fallback Chain
~~~~~~~~~~~~~~

Strategy selection and fallback logic when the primary solver fails.

.. automodule:: homodyne.optimization.nlsq.fallback_chain
   :members:
   :undoc-members:
   :show-inheritance:

Error Recovery
~~~~~~~~~~~~~~

3-attempt error recovery with progressive conservative settings and
5-category error diagnosis.

.. automodule:: homodyne.optimization.nlsq.recovery
   :members:
   :undoc-members:
   :show-inheritance:

Stratified Least Squares
~~~~~~~~~~~~~~~~~~~~~~~~

Primary fitting strategy using stratified JTJ accumulation via NLSQ
``least_squares()``.

.. automodule:: homodyne.optimization.nlsq.strategies.stratified_ls
   :members:
   :undoc-members:
   :show-inheritance:

Hybrid Streaming
~~~~~~~~~~~~~~~~

4-phase adaptive hybrid streaming optimizer for large datasets.

.. automodule:: homodyne.optimization.nlsq.strategies.hybrid_streaming
   :members:
   :undoc-members:
   :show-inheritance:

Out-of-Core Accumulation
~~~~~~~~~~~~~~~~~~~~~~~~

Global JTJ accumulation solver for datasets exceeding the memory threshold.

.. automodule:: homodyne.optimization.nlsq.strategies.out_of_core
   :members:
   :undoc-members:
   :show-inheritance:
