# Timing Instrumentation for Reflectorch Predictions

## Overview

The [`TimedInferenceModel`](timed_inference_model.py) extends the standard
`reflectorch.InferenceModel` to provide detailed timing measurements for
reflectometry predictions. This allows tracking of:

- **NN Prediction Time**: Time spent on neural network inference
- **Fitting Time**: Time spent on scipy optimization (polishing)
- **Total Time**: Combined execution time

## Architecture

```
InferenceModel (reflectorch)
    ↓
TimedInferenceModel (vipr_reflectometry/reflectorch/load_model/)
    ↓
Used by ReflectorchModelLoader
    ↓
Timing results in prediction output
    ↓
ReflectorchPerformance hook → stored in batch_metadata['prediction_timing']
```

## How It Works

`TimedInferenceModel` overrides two key methods:

1. **`predict()`**: Starts timing when called, adds timing results to output.
2. **`_polish_prediction()`**: Captures the transition point between NN
   prediction and fitting.

This approach uses **inheritance** rather than monkey-patching, making it:
- Clean and maintainable
- Type-safe
- Easy to understand and debug

## Usage

The timing feature is **automatically enabled** for all reflectorch predictions.
No configuration is needed.

### Single Spectrum Example

```python
from vipr_reflectometry.reflectorch.load_model import TimedInferenceModel

model = TimedInferenceModel(config_name='mc66', device='cpu')

result = model.predict(
    reflectivity_curve=curve,
    q_values=q_values,
    prior_bounds=bounds,
    polish_prediction=True
)

timing = result['timing']
print(f"NN: {timing['nn_prediction_s']:.3f}s")
print(f"Fitting: {timing['fitting_s']:.3f}s")
print(f"Total: {timing['total_s']:.3f}s")
```

### Output Format

#### Single Spectrum

```python
{
    'predicted_params_array': [...],
    'polished_params_array': [...],
    'timing': {
        'nn_prediction_s': 0.123,    # NN inference time
        'fitting_s': 2.456,          # Scipy fitting time
        'total_s': 2.579             # Total time
    }
}
```

#### Batch Processing

For batch predictions, timing statistics are aggregated per spectrum:

```python
{
    'predicted_params_array': [...],  # shape: (batch_size, n_params)
    'polished_params_array': [...],
    'timing': {
        'nn_prediction_s': {
            'mean': 0.125,
            'std': 0.012,
            'min': 0.110,
            'max': 0.145
        },
        'fitting_s': {
            'mean': 2.450,
            'std': 0.320,
            'min': 2.100,
            'max': 3.200
        },
        'total_s': {
            'mean': 2.575,
            'std': 0.328,
            'min': 2.210,
            'max': 3.345
        }
    }
}
```

## Logging

Timing information is automatically logged:

### Single Spectrum

```
[INFO] Prediction timing: NN=0.123s, Fitting=2.456s, Total=2.579s
```

### Batch Processing

```
[INFO] Batch timing averages: NN=0.125s (±0.012s), Fitting=2.450s (±0.320s), Total=2.575s (±0.328s)
```

## Performance Overhead

The timing instrumentation has minimal overhead:

- `time.perf_counter()` overhead: ~100 ns per call
- Total overhead per prediction: < 1 ms
- Memory overhead: ~48 bytes per spectrum (3 `float64` values)

## Persistence

The [`ReflectorchPerformance`](../performance.py) hook runs at
`INFERENCE_COMPLETE_HOOK` and stores the formatted timing data in:

```python
app.datacollector.data.batch_metadata['prediction_timing']
```

This means timing data is persisted to `data.pkl` alongside all other result
data and is accessible for UI rendering and export.

## Files

| File | Purpose |
|---|---|
| [`timed_inference_model.py`](timed_inference_model.py) | `InferenceModel` subclass with timing |
| [`reflectorch_model_loader.py`](reflectorch_model_loader.py) | Uses `TimedInferenceModel` instead of `EasyInferenceModel` |
| [`../performance.py`](../performance.py) | Hook-based persistence of timing data |
| [`../predict/reflectorch_predictor.py`](../predict/reflectorch_predictor.py) | Logging + batch aggregation |

## Disabling Timing

To disable timing, revert `reflectorch_model_loader.py` to use
`EasyInferenceModel` directly:

```python
from reflectorch import EasyInferenceModel
...
model = EasyInferenceModel(config_name=config_name, ...)
```

## Troubleshooting

### `timing` key missing in results

- Ensure `TimedInferenceModel` is used (check `reflectorch_model_loader.py`)
- Verify the model was loaded correctly (check logs for initialization errors)

### Unexpected timing values

- NN prediction typically takes **0.1 – 0.5 s**
- Fitting typically takes **1 – 10 s** depending on data complexity
- Check system load, device (CPU vs GPU), and input data size if values differ
  significantly from these ranges
