nlsq.adaptive\_hybrid\_streaming module
=======================================

.. currentmodule:: nlsq.streaming.adaptive_hybrid

.. automodule:: nlsq.streaming.adaptive_hybrid
   :noindex:

Overview
--------

The ``nlsq.adaptive_hybrid_streaming`` module implements a four-phase hybrid
optimizer that solves three fundamental issues in streaming optimization:

1. **Weak gradient signals** from parameter scale imbalance (via normalization)
2. **Slow convergence** near optimum (via Gauss-Newton)
3. **Crude covariance estimation** (via exact J^T J accumulation)

**New in version 0.3.0**: Complete adaptive hybrid streaming optimizer.

Key Features
------------

- **Four-phase optimization**: Automatic phase transitions for optimal convergence
- **Parameter normalization**: Address gradient signal weakness from scale imbalance
- **L-BFGS warmup**: Quasi-Newton optimization with 4-layer divergence protection
- **Streaming Gauss-Newton**: Second-order convergence with streaming J^T J
- **Exact covariance**: Production-quality uncertainty estimates
- **Fault tolerance**: Checkpointing, validation, and automatic recovery
- **Multi-device support**: GPU/TPU parallelism for large datasets
- **Defense telemetry**: Production monitoring of warmup protection layers

**New in version 0.3.6**: 4-Layer Defense Strategy for warmup divergence prevention.

Optimization Phases
-------------------

Phase 0: Normalization Setup
~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Sets up parameter normalization to address gradient signal weakness:

- Determines normalization strategy (bounds-based, p0-based, or none)
- Creates ``ParameterNormalizer`` with scales and offsets
- Wraps user model for transparent normalized parameter space
- Transforms bounds to normalized space
- Stores normalization Jacobian for Phase 3

Phase 1: L-BFGS Warmup with 4-Layer Defense
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

First-order optimization with adaptive switching and divergence protection:

- Uses Optax L-BFGS optimizer with line search and configurable step sizes
- **4-Layer Defense Strategy** (new in 0.3.6):

  1. **Warm Start Detection**: Skip warmup if initial loss already low
  2. **Adaptive Step Size**: Scale step size based on initial loss quality
  3. **Cost-Increase Guard**: Abort if loss increases beyond tolerance
  4. **Step Clipping**: Limit update magnitude for stability

- Monitors loss plateau and gradient norm for switch criteria
- Builds momentum and explores parameter space
- Switches to Phase 2 when ready for fine-tuning

.. seealso::

   :doc:`../explanation/how_fitting_works` for complete optimization strategy documentation.

Phase 2: Streaming Gauss-Newton
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Second-order optimization with exact J^T J accumulation:

- Streams data in chunks for memory efficiency
- Accumulates exact J^T J matrix (not stochastic approximation)
- Trust region step control with Levenberg-Marquardt regularization
- Converges quickly near optimum with quadratic rate
- Provides production-quality parameter estimates

Phase 3: Denormalization and Covariance
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Finalizes optimization and computes uncertainties:

- Denormalizes parameters to original space
- Transforms covariance matrix using normalization Jacobian
- Returns final result with parameter uncertainties

Classes
-------

.. autosummary::
   :toctree: generated/
   :template: class.rst

   AdaptiveHybridStreamingOptimizer
   DefenseLayerTelemetry

.. autoclass:: AdaptiveHybridStreamingOptimizer
   :members:
   :undoc-members:
   :show-inheritance:
   :noindex:

.. autoclass:: DefenseLayerTelemetry
   :members:
   :undoc-members:
   :show-inheritance:
   :noindex:

Functions
---------

.. autosummary::
   :toctree: generated/

   get_defense_telemetry
   reset_defense_telemetry

.. autofunction:: get_defense_telemetry
   :noindex:

.. autofunction:: reset_defense_telemetry
   :noindex:

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

Basic Usage
~~~~~~~~~~~

Create an optimizer with default configuration:

.. code-block:: python

    from nlsq import AdaptiveHybridStreamingOptimizer, HybridStreamingConfig
    import jax.numpy as jnp

    # Default configuration
    config = HybridStreamingConfig()
    optimizer = AdaptiveHybridStreamingOptimizer(config)


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


    # Prepare data
    x_data = jnp.linspace(0, 10, 1000000)  # 1M points
    y_data = 2.5 * jnp.exp(-0.3 * x_data) + 0.5 + 0.1 * jnp.random.normal(key, (1000000,))

    # Fit (API for full implementation)
    result = optimizer.fit(
        model=exponential,
        x_data=x_data,
        y_data=y_data,
        p0=[2.0, 0.5, 0.3],
    )

With Bounds-Based Normalization
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Use parameter bounds for normalization:

.. code-block:: python

    config = HybridStreamingConfig(
        normalize=True,
        normalization_strategy="bounds",
    )
    optimizer = AdaptiveHybridStreamingOptimizer(config)

    # Parameters: amplitude [1, 10], decay [0.1, 1.0], offset [0, 2]
    bounds = (jnp.array([1.0, 0.1, 0.0]), jnp.array([10.0, 1.0, 2.0]))

    result = optimizer.fit(
        model=exponential,
        x_data=x_data,
        y_data=y_data,
        p0=[2.0, 0.5, 0.3],
        bounds=bounds,
    )

With Aggressive Configuration
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Fast convergence with looser tolerances:

.. code-block:: python

    config = HybridStreamingConfig.aggressive()
    optimizer = AdaptiveHybridStreamingOptimizer(config)

    # Faster warmup, larger chunks, earlier switching
    result = optimizer.fit(model, x_data, y_data, p0)

With Conservative Configuration
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Higher quality with tighter tolerances:

.. code-block:: python

    config = HybridStreamingConfig.conservative()
    optimizer = AdaptiveHybridStreamingOptimizer(config)

    # Tighter tolerance, more Gauss-Newton iterations
    result = optimizer.fit(model, x_data, y_data, p0)

Memory-Optimized Configuration
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Minimize memory footprint:

.. code-block:: python

    config = HybridStreamingConfig.memory_optimized()
    optimizer = AdaptiveHybridStreamingOptimizer(config)

    # Smaller chunks, float32, frequent checkpoints
    result = optimizer.fit(model, x_data, y_data, p0)

Custom Switching Criteria
~~~~~~~~~~~~~~~~~~~~~~~~~

Control when Phase 1 switches to Phase 2:

.. code-block:: python

    config = HybridStreamingConfig(
        warmup_iterations=300,
        max_warmup_iterations=800,
        loss_plateau_threshold=1e-5,  # Tighter plateau detection
        gradient_norm_threshold=1e-4,  # Lower gradient threshold
        active_switching_criteria=["plateau", "gradient"],  # Remove max_iter
    )
    optimizer = AdaptiveHybridStreamingOptimizer(config)

With Fault Tolerance
~~~~~~~~~~~~~~~~~~~~

Enable checkpointing for long optimizations:

.. code-block:: python

    config = HybridStreamingConfig(
        enable_checkpoints=True,
        checkpoint_frequency=50,
        checkpoint_dir="/path/to/checkpoints",
        validate_numerics=True,
        enable_fault_tolerance=True,
        max_retries_per_batch=3,
    )
    optimizer = AdaptiveHybridStreamingOptimizer(config)

    # Resume from checkpoint
    config_resume = HybridStreamingConfig(resume_from_checkpoint="/path/to/checkpoint.pkl")

Phase Tracking
~~~~~~~~~~~~~~

Monitor optimization progress:

.. code-block:: python

    optimizer = AdaptiveHybridStreamingOptimizer(config)

    # After fitting
    print(f"Final phase: {optimizer.current_phase}")
    print(f"Phase history: {optimizer.phase_history}")

    # Each entry in phase_history:
    # {'phase': 0, 'start_time': 1234567890.0, 'duration': 0.5}

Defense Layer Telemetry
~~~~~~~~~~~~~~~~~~~~~~~

**New in version 0.3.6**: Monitor defense layer activations.

Track when defense layers trigger across multiple fits:

.. code-block:: python

    from nlsq import get_defense_telemetry, reset_defense_telemetry

    # Reset telemetry counters
    reset_defense_telemetry()

    # Run multiple fits
    for dataset in datasets:
        result = optimizer.fit(model, x_data, y_data, p0)

    # Get telemetry summary
    telemetry = get_defense_telemetry()
    print(telemetry.get_summary())

    # Get activation rates
    rates = telemetry.get_trigger_rates()
    print(f"Warm start rate: {rates['layer1_warm_start_rate']:.1f}%")
    print(f"Cost guard rate: {rates['layer3_cost_guard_rate']:.1f}%")

    # Export for Prometheus/Grafana
    metrics = telemetry.export_metrics()

Defense Layer Presets
~~~~~~~~~~~~~~~~~~~~~

**New in version 0.3.6**: Sensitivity presets for defense layers.

.. code-block:: python

    # Strict defense for near-optimal scenarios
    config = HybridStreamingConfig.defense_strict()

    # Relaxed defense for exploration
    config = HybridStreamingConfig.defense_relaxed()

    # Disable all defense layers
    config = HybridStreamingConfig.defense_disabled()

    # Scientific computing optimized
    config = HybridStreamingConfig.scientific_default()

Architecture
------------

Optimization Flow
~~~~~~~~~~~~~~~~~

.. code-block:: text

    Input: model, x_data, y_data, p0, bounds
           │
           ▼
    ┌──────────────────────────────────────┐
    │  Phase 0: Normalization Setup        │
    │  - Create ParameterNormalizer        │
    │  - Wrap model function               │
    │  - Transform bounds                  │
    │  - Store normalization Jacobian      │
    └──────────────────────────────────────┘
           │
           ▼
    ┌──────────────────────────────────────┐
    │  Phase 1: L-BFGS Warmup              │
    │  - Optax L-BFGS optimizer            │
    │  - Monitor loss/gradient             │
    │  - Check switching criteria          │
    └──────────────────────────────────────┘
           │ (switch when criteria met)
           ▼
    ┌──────────────────────────────────────┐
    │  Phase 2: Streaming Gauss-Newton     │
    │  - Stream data in chunks             │
    │  - Accumulate exact J^T J            │
    │  - Trust region optimization         │
    └──────────────────────────────────────┘
           │ (converged)
           ▼
    ┌──────────────────────────────────────┐
    │  Phase 3: Denormalization            │
    │  - Denormalize parameters            │
    │  - Transform covariance: J @ C @ J.T │
    │  - Return final result               │
    └──────────────────────────────────────┘
           │
           ▼
    Output: OptimizeResult with popt, pcov

Key Attributes
~~~~~~~~~~~~~~

.. list-table::
   :widths: 25 75
   :header-rows: 1

   * - Attribute
     - Description
   * - ``current_phase``
     - Current optimization phase (0, 1, 2, or 3)
   * - ``phase_history``
     - List of phase transitions with timing
   * - ``normalizer``
     - ParameterNormalizer instance
   * - ``normalized_model``
     - Wrapped model for normalized space
   * - ``normalized_bounds``
     - Bounds in normalized space
   * - ``normalization_jacobian``
     - Jacobian for covariance transform
   * - ``best_params_global``
     - Best parameters found (for fault tolerance)
   * - ``best_cost_global``
     - Best cost found (for fault tolerance)

Performance Characteristics
---------------------------

Convergence Speed
~~~~~~~~~~~~~~~~~

- **Phase 1 (L-BFGS)**: Superlinear convergence, robust to initialization
- **Phase 2 (Gauss-Newton)**: Quadratic convergence near optimum
- **Overall**: Faster than pure first-order or pure second-order

Memory Usage
~~~~~~~~~~~~

- **Streaming**: Fixed memory footprint regardless of dataset size
- **J^T J accumulation**: O(n_params^2) memory
- **Chunks**: Configurable via ``chunk_size``

Throughput
~~~~~~~~~~

- CPU: 50,000 - 200,000 samples/second
- GPU: 500,000 - 2,000,000 samples/second
- Depends on model complexity

When to Use
-----------

**Use Adaptive Hybrid Streaming when:**

- Parameters span many orders of magnitude (gradient imbalance)
- Dataset is large (100K+ points)
- Need production-quality uncertainty estimates
- Standard optimizers converge slowly
- Memory is limited relative to dataset size

**Use standard curve_fit when:**

- Dataset fits in memory
- Parameters have similar scales
- Simple models with good initialization
- Don't need streaming

See Also
--------

- :doc:`nlsq.hybrid_streaming_config` : Configuration options
- :doc:`nlsq.parameter_normalizer` : Parameter normalization
- :doc:`nlsq.stability` : Numerical stability utilities
- :doc:`../howto/handle_large_data` : Large dataset guide
