feat: AI slope-completion fallback for degraded SH-WFS frames (LSTM, temporal consistency) #94

Summary

Part of umbrella feature request #90 (AI and Radio fallback modes). Implement a slope-domain sequence model that reconstructs missing or corrupted wavefront slopes (Δx, Δy) per sub-aperture from the temporal consistency of past frames, triggered by the frame-quality gate (#93) when the detector reports partial signal loss (mild fog, scattering, or partial occlusion). The completed slopes are fed to the existing zonal/modal reconstructors so the loop continues without forcing the DM into impossible geometries.

1. Background and Motivation

• The existing ML stack predicts future Zernike coefficients (20 modes) via LSTM (sequence_models.py:5-32, predictive_ao.py:23-31, train_sequence.py:78-81). The proposed route needs current missing slopes — a different input/output schema and a different objective (signal recovery vs. latency compensation).
• Reusable components: LSTM/GRU architecture, sequence training pipeline (SHSequenceDataset, split_by_sequence, train_sequence.py:30-128), ONNX export (export_onnx.py), C ring-buffer state machinery (predictive_ao.c).
• The C ONNX Runtime integration is currently a disabled #if 0 stub (predictive_ao.c:16-93); the active path returns NULL/-1. This must be completed for real-time use.
• Native CUDA ML inference kernels do not exist; the GPU ML figures in docs/paper are PyTorch-based, not C/CUDA.
• Current degraded-frame behavior is spatial-only nearest-neighbor interpolation (stream.c:277-318), which ignores the strong temporal correlation of atmospheric turbulence (Taylor frozen flow) — the opportunity this issue exploits.

2. Problem Formulation

2.1 Data model

Let the per-spot slope vector at frame t be

s_t = [Δx_{1,t}, Δy_{1,t}, …, Δx_{N,t}, Δy_{N,t}]ᵀ ∈ ℝ^{2N}

where N = cal->nspots (137 for synthetic data). The observation mask is m_t ∈ {0,1}^{2N} with m_t[k] = 1 for a valid spot and 0 for a lost/corrupted one (already produced by rippa_compute_deltas). We observe only the masked subset

s̃_t = m_t ⊙ s_t   (⊙ = elementwise/Hadamard product)

2.2 Latent state-space model

Model slope dynamics as a latent linear-Gaussian AR(1) process (consistent with the project's Kolmogorov + Taylor frozen-flow AR(1) generator in synthetic_shwfs.py):

z_t = A z_{t-1} + η_t,    η_t ~ N(0, Q)          (process noise)
s_t = C z_t + ε_t,        ε_t ~ N(0, R)          (measurement noise)

with latent state z_t ∈ ℝ^d, transition A, emission C, noise covariances Q, R.

2.3 Objective

Given L frames of history and the partial current observation, we seek the posterior mode of the full current slope vector:

ŝ_t = argmax_s p(s_t | s̃_{t-L:t-1}, s̃_t, m_{t-L:t})

This is a filtering/smoothing problem; under linear-Gaussian assumptions the exact solution is the Kalman filter. However, real turbulence is nonlinear in the latent dynamics and the masking pattern is structured (spatially contiguous fog banks), so we use a neural amortized posterior instead.

2.4 Network as amortized posterior

A recurrent sequence encoder (LSTM or GRU) with hidden state h_t amortizes the filtering distribution:

h_t = f_θ(h_{t-1}, s̃_t, m_t)
ŝ_t = g_θ(h_t, m_t)

• Inputs (per time step): zero-filled observed slopes s̃_t (shape [2N]) concatenated with the mask m_t (shape [2N]), total [2N + 2N] = [4N].
• Context window: L past frames + the current partial frame → tensor [L+1, 4N].
• Output: ŝ_t ∈ ℝ^{2N} — the complete slope vector for the current frame (masked entries reconstructed, valid entries optionally passed through).
• Mask-conditioning is essential: the network must learn where it is allowed to be uncertain and propagate that into the recurrence.

3. Loss Function

3.1 Masked reconstruction loss (primary)

Train only on the missing/corrupted entries, so the model learns completion rather than reproduction of known data:

L_rec(θ) = ( 1 / Σ_k (1 − m_t[k]) ) · Σ_k (1 − m_t[k]) · (ŝ_t[k] − s_t[k])²

Normalized by the number of masked entries so the loss is scale-invariant to mask density.

3.2 Zernike-consistency regularizer

Enforce physical plausibility by projecting the prediction onto the Zernike derivative basis (the modal model Z′ already built in rippra_modal_setup, recon.c):

a_t = (Z′)⁺ ŝ_t                (pseudo-inverse, rcond = 1e-4 as in recon.c:359)
ŝ_t^Z = Z′ a_t
L_zern = ||ŝ_t − ŝ_t^Z||²

This penalizes slope fields that cannot arise from a smooth Zernike expansion, directly addressing the actuator-command-space concern raised in #90.

3.3 Temporal smoothness regularizer

Encourage continuity with the previous frame (Taylor frozen-flow assumption):

L_tem = ||ŝ_t − s_{t-1}||²

3.4 Total objective

L(θ) = L_rec + λ_zern · L_zern + λ_tem · L_tem

with hyperparameters λ_zern, λ_tem tuned on validation (suggested starting points: λ_zern = 1e-2, λ_tem = 1e-1).

4. Physical Constraint Enforcement (Safety)

To guarantee the model output maps into valid actuator space (the core concern in #90):

1. Zernike projection: replace the raw network output with its projection ŝ_t^Z computed above (also truncates unstable high-order modes via the rcond cutoff).
2. Stroke clamping: per-spot limit

   |Δx_t[k]| ≤ pitch_px/2,   |Δy_t[k]| ≤ pitch_px/2
   

   where pitch_px ≈ 40.5 px (config pitch/camera_pixsize), i.e. the cross-talk limit beyond which a spot bleeds into a neighboring sub-aperture.
3. DM command bound (optional, coordinate with DM-safety track): the reconstructed phase → DM map should never exceed the configured max stroke; if the DM-safety issue lands first, this model simply guarantees its input is sane.

The constrained output ŝ_t^constrained is what reaches rippra_zonal_reconstruct / rippra_modal_reconstruct.

5. Implementation Plan

5.1 Dataset generation (synthetic_shwfs.py)

• Extend the generator to emit masked slope sequences: simulate partial attenuation by scaling spot intensities in contiguous spatial blocks (structured fog) and/or random Bernoulli spot dropout, producing (s_t, m_t) pairs.
• Produce a masking schedule covering: clean, mild fade (10-40% spots), moderate fog blocks (40-70%), and heavy (70-100%) — the latter being the boundary with the hard fallback.
• Keep the existing physical model (Kolmogorov AR(1), Zernike → slopes) unchanged; only add the observation/masking layer.

5.2 Model (sequence_models.py or models.py)

• New class SlopeCompletionLSTM: input [batch, L+1, 4N], hidden size 128, num_layers = 2, output head → [batch, 2N].
• Mask-conditioned (mask tensor is a model input, not just a training artifact).

5.3 Training (train_sequence.py)

• Add a slope-completion task alongside existing predict/classify/param-est tasks.
• Sequence-level split with split_by_sequence() and the check_split_leakage() assertion (already enforced for other tasks) to prevent leakage.
• Targets = original unmasked slopes; loss per §3.

5.4 Export (export_onnx.py)

• Export with dynamic batch axis; dummy input [1, L+1, 4N].
• Verify with test_onnx_models.py (round-trip parity).

5.5 C integration (predictive_ao.c)

• Complete the ONNX Runtime path: enable the #if 0 block (predictive_ao.c:16-93), link onnxruntime (CPU EP by default, CUDA EP where available).
• Add API:

  int predictive_ao_complete_slopes(PredictiveAOState *state,
                                    const double *observed_dx,
                                    const double *observed_dy,
                                    const int *mask, int nspots,
                                    double *out_dx, double *out_dy);

  The state's ring buffer holds the last L slope frames (extend the existing 10-frame Zernike history to slopes, or add a parallel buffer).
• Return codes: 0 = success, 1 = model unavailable (graceful fallback to spatial interpolation), negative = error.

5.6 Stream wiring (stream.c)

• Register the completion hook via the feat: add frame-quality gate (Dynamic Signal Routing) before wavefront reconstruction #93 fallback callback (rippra_stream_set_fallback).
• On status = DEGRADED: call predictive_ao_complete_slopes, then run reconstruction with the completed slopes.
• On hook failure or status = LOST: existing spatial interpolation / hard fallback respectively.

5.7 Benchmarking

• Latency budget: inference must be < 1 ms (vs. the ~10 ms τ₀ coherence budget; paper reports classical path 0.041 ms, full pipeline ~0.9 ms).
• Use benchmark_onnx_latency.py pattern; report CPU vs CUDA EP.

6. Mathematical Notes / Validation

• Baseline: compare against the existing spatial nearest-neighbor interpolation (stream.c:277-318) on the same masked sequences.
• Target: masked-completion RMSE ≤ 30% of the spatial-interpolation baseline on held-out sequences, at moderate mask densities (≤ 50%).
• Verify temporal consistency benefit: for a fixed mask pattern, RMSE should decrease as L (context length) grows up to a saturation point (expected ~10 frames per existing LSTM findings).
• Confirm the Zernike-projection regularizer reduces out-of-range actuator commands (adversarial mask test).
• Validate on the existing sh_flat.raw/img.raw pipeline: a synthetically masked img.raw should be reconstructed to within a tolerance of the unmasked reconstruction.

7. Acceptance Criteria

• SlopeCompletionLSTM trained with mask-weighted loss; masked-completion RMSE ≤ 30% of spatial-interpolation baseline
• Sequence-level split + check_split_leakage assertion passes
• Masked-loss verified: gradients only from missing entries
• ONNX export round-trips correctly (test_onnx_models.py green)
• predictive_ao_complete_slopes functional with ONNX Runtime; graceful fallback when model unavailable
• Inference latency < 1 ms (CPU EP; CUDA EP where available)
• Physical constraints (Zernike projection + stroke clamp) verified with adversarial masks
• Stream integration via feat: add frame-quality gate (Dynamic Signal Routing) before wavefront reconstruction #93 hook: DEGRADED frames use completed slopes; LOST frames park DM
• Existing 38/38 C tests and Python ML CI jobs still pass
• Docs: docs/algorithms.md, docs/paper/rippra_paper.tex (brief mention), ML README updated

8. Dependencies / Relationship to Other Issues

• Blocked on: feat: add frame-quality gate (Dynamic Signal Routing) before wavefront reconstruction #93 (frame-quality gate) for stream integration; model + dataset work can proceed in parallel.
• Reuses: LSTM architecture, SHSequenceDataset, ONNX toolchain, predictive_ao.c machinery, synthetic_shwfs.py.
• Related: DM safety/clamping track (complementary safety layer); Feature Request: Add AI and Radio fallback modes to prevent mirror crashes during bad weather #90 (umbrella).
• Extends: sequence_models.py, train_sequence.py, synthetic_shwfs.py, export_onnx.py, predictive_ao.c, stream.c.

9. Out of Scope

• Hard fallback (RF handshake, low-power polling) — hardware, Feature Request: Add AI and Radio fallback modes to prevent mirror crashes during bad weather #90
• Native CUDA ML inference kernels — separate optimization track
• DM clamping/park policy — separate track
• Future-Zernike predictive control (existing capability, unchanged by this issue)