# Metrics and Recording

Deep-dive into `TorchMetricRecorder`, metric builders, and how metrics flow through the training loop.

## TorchMetricRecorder

Source: `metric/torch_metric.py:158-429`

The central metric recording component. Combines two types of metrics:
- **Computed metrics**: torchmetrics `Metric` / `MetricCollection` objects (Accuracy, F1, etc.)
- **Scalar metrics**: plain float values from step function output (loss, custom losses)

### Constructor

```python
TorchMetricRecorder(
    metrics: MetricCollection | dict[str, Metric] | None = None,
    before_metric: 'sigmoid' | 'softmax' | Callable | dict[str, Callable] | None = None,
    scalar_metrics: ScalarMetric | str | list[str] | Literal['auto'] | None = 'loss',
)
```

### Parameter: metrics

Torchmetrics-based evaluation metrics. Built using `clf_metrics_builder()` or `reg_metrics_builder()`.

```python
# Binary classification
recorder = TorchMetricRecorder(clf_metrics_builder(binary=True))

# Multiclass
recorder = TorchMetricRecorder(clf_metrics_builder(task='multiclass', num_classes=10))

# Regression
recorder = TorchMetricRecorder(reg_metrics_builder())

# Custom metrics
from torchmetrics.classification import Accuracy, AUROC
recorder = TorchMetricRecorder({'acc': Accuracy(task='binary'), 'auroc': AUROC(task='binary')})

# No computed metrics (only scalar tracking)
recorder = TorchMetricRecorder(None, scalar_metrics='auto')
```

### Parameter: before_metric

Activation function applied to model predictions **before** passing to torchmetrics. This is critical — wrong activation produces garbage metric values.

| Mode | Value | Behavior |
|------|-------|----------|
| None | `None` | No activation. Use when model output is already activated. |
| String preset | `'sigmoid'` | `torch.sigmoid(y_pred)`. For binary clf with `BCEWithLogitsLoss`. |
| String preset | `'softmax'` | `torch.softmax(y_pred, dim=-1)`. For multiclass with `CrossEntropyLoss`. |
| Callable | `lambda x: x.argmax(-1)` | Any `(Tensor) -> Tensor` function. |
| Dict (glob) | `{'accuracy': torch.sigmoid, 'U*': None}` | Per-metric activation with glob pattern matching. |

**Dict-glob matching rules** (source: `torch_metric.py:260-310`):
1. Exact match first: `'accuracy'` matches metric named `accuracy`
2. Glob pattern: `'U*'` matches `Uncertainty_aleatoric`, `Uncertainty_epistemic`, etc.
3. Wildcard: `'*'` matches all metrics
4. No match: uses raw predictions (no activation)
5. Results are cached after first match

### Parameter: scalar_metrics

Controls which float values from step function output are tracked.

| Mode | Value | Behavior |
|------|-------|----------|
| String | `'loss'` | Track single metric named `loss` (default) |
| List | `['loss', 'ce_loss', 'reg_loss']` | Track multiple named metrics |
| `'auto'` | `'auto'` | Auto-define: any key passed to `update()` is automatically tracked |
| `ScalarMetric` | instance | Use a pre-configured `ScalarMetric` object |
| `None` | `None` | No scalar tracking (only torchmetrics) |

**Important**: If step function returns `{'g_loss': 0.5, 'd_loss': 0.3}` but `scalar_metrics='loss'`, both `g_loss` and `d_loss` are silently ignored. Use `scalar_metrics='auto'` or `['g_loss', 'd_loss']`.

### update() and compute()

```python
# Called per-batch by train_epoch:
metrics = recorder.update(y_pred, y_target, **raw_metrics_from_step)
# raw_metrics_from_step = {'loss': 0.5} from training_step return value

# Called per-epoch by train_epoch:
epoch_metrics = recorder.compute()  # returns averaged metrics, auto-resets
```

`compute()` automatically resets both torchmetrics and scalar metrics for the next epoch.

### Device handling

Torchmetrics must be on the same device as tensors. `ctx.launch()` handles this:
```python
ctx.launch()  # calls fabric.launch() + recorder.to_device(fabric.device)
```

If you skip `ctx.launch()` and use `fabric.launch()` directly, you must manually call:
```python
train_recorder.to_device(fabric.device)
val_recorder.to_device(fabric.device)
```

## clf_metrics_builder

Source: `metric/__init__.py:10-127`

Factory function for classification metric collections.

```python
clf_metrics_builder(
    binary: bool = True,
    *,
    task: Literal['binary', 'multiclass', 'multilabel'] | None = None,
    threshold: float = 0.5,
    num_classes: int | None = None,
    average: Literal['micro', 'macro', 'weighted', 'none'] | None = 'macro',
    top_k: int | None = 1,
    use_auc: bool = False,
    use_specificity: bool = False,
    use_matthews: bool = False,
    use_kappa: bool = False,
    use_ece: bool = False,
) -> MetricCollection
```

**Base metrics** (always included): `accuracy`, `precision`, `recall`, `f1`

**Optional metrics** (via boolean flags):
| Flag | Metrics added | Notes |
|------|--------------|-------|
| `use_auc=True` | `auroc`, `auprc` | Binary/multiclass only |
| `use_specificity=True` | `specificity` | Binary/multiclass only |
| `use_matthews=True` | `matthews_corrcoef` | All tasks |
| `use_kappa=True` | `cohen_kappa` | Binary/multiclass only |
| `use_ece=True` | `ece` (CalibrationError) | All tasks |

**Task resolution**: If `task` is None, uses `binary` param: `True` -> `'binary'`, `False` -> `'multiclass'`.

**Multiclass requirement**: `num_classes` must be provided for `'multiclass'` and `'multilabel'` tasks.

## reg_metrics_builder

Source: `metric/__init__.py:129-143`

Factory function for regression metric collections. No parameters.

```python
reg_metrics_builder() -> MetricCollection
```

Returns: `mae`, `mse`, `rmse`, `r2_score`

## ScalarMetric

Source: `metric/torch_metric.py:19-156`

Low-level scalar tracking. Usually you don't create this directly — `TorchMetricRecorder` creates one internally based on `scalar_metrics` parameter.

```python
scalar = ScalarMetric(default_reducer='mean', auto_define=True)
scalar.update(loss=0.5, custom_metric=0.8)  # auto-defines on first use
result = scalar.compute()  # {'loss': 0.5, 'custom_metric': 0.8}
scalar.reset()
```

Modes:
- `auto_define=False` (default): must call `define_metric('name', 'mean')` before use
- `auto_define=True`: automatically defines metrics on first `update()` call

## Custom metrics with create_lambda_metric

Source: `metric/common_metric.py`

For domain-specific metrics not available in torchmetrics (e.g. signal-to-noise ratio, spectral similarity), use `create_lambda_metric` to wrap any `(preds, target) -> Tensor` function into a torchmetrics-compatible `Metric`:

```python
from zyplib.train.metric import create_lambda_metric

def my_snr(preds, target):
    signal_power = torch.mean(target ** 2)
    noise_power = torch.mean((preds - target) ** 2)
    return 10 * torch.log10(signal_power / (noise_power + 1e-8))

snr_metric = create_lambda_metric('SNR', my_snr, higher_is_better=True, reduce_fn='mean')

# Use in TorchMetricRecorder alongside standard metrics
recorder = TorchMetricRecorder(
    {'accuracy': Accuracy(task='binary'), 'snr': snr_metric},
    before_metric='sigmoid',
)
```

## Common configurations

| Task | loss_fn | before_metric | clf_metrics_builder |
|------|---------|---------------|---------------------|
| Binary clf | `BCEWithLogitsLoss()` | `'sigmoid'` | `clf_metrics_builder(binary=True)` |
| Multiclass clf | `CrossEntropyLoss()` | `'softmax'` | `clf_metrics_builder(task='multiclass', num_classes=N)` |
| Regression | `MSELoss()` | `None` | `reg_metrics_builder()` |
| GAN | custom | `None` | `None` (scalar_metrics only) |
| EDL | model.compute_loss | `None` | `clf_metrics_builder(...)` + custom scalar_metrics |
