# Paradigm Extensions

How to extend `zyplib.train` for non-standard training paradigms. All extensions follow the same pattern.

## The general pattern

1. Write a **Procedure class** with `xxx_train_step` and `xxx_val_step` methods
2. These methods must follow the **step function contract** (see `step-function-contract.md`)
3. Wire them into `procedure.train_loop` via `fn_train_step=` / `fn_val_step=`
4. Configure `scalar_metrics` to match the custom loss names returned by your step

```
Standard:  training_step -> returns (y_pred, {'loss': float})
Custom:    my_train_step -> returns (y_pred, {'loss': float, 'extra_loss': float, ...})
                                     ^                ^
                                     |                |
                          TorchMetricRecorder    ScalarMetric must know these keys
```

## GAN: BasicGANFabricProcedure

Source: `zyplib/paradigm/gan.py:219-281`

GAN training has unique requirements:
- **Two optimizers** (generator + discriminator) stored as a tuple in `ctx.optimizer`
- **No standard labels** — batch is `(real_data,)` or `(real_data, _)`
- **Custom loss names**: `g_loss`, `d_loss` instead of `loss`
- **y_pred is None** — GAN doesn't produce classification predictions

### Setup

```python
from zyplib.paradigm.gan import StandardGAN, BasicGANFabricProcedure
from zyplib.train import default_fabric, FabricTrainingContext, procedure, TorchMetricRecorder

# GAN model wraps generator + discriminator
gan_model = StandardGAN(generator, discriminator, latent_dim=100, real_data_shape=(1, 28, 28))

# Two optimizers as a tuple
optimizer_g = torch.optim.Adam(generator.parameters(), lr=2e-4, betas=(0.5, 0.999))
optimizer_d = torch.optim.Adam(discriminator.parameters(), lr=2e-4, betas=(0.5, 0.999))

fabric = default_fabric(
    dir_name='gan_logs',
    enabled={'save_checkpoints_cb': False},  # GAN checkpointing needs special handling
)

# No torchmetrics — only scalar tracking for g_loss and d_loss
recorder = TorchMetricRecorder(
    metrics=None,
    scalar_metrics=['g_loss', 'd_loss'],
)

ctx = FabricTrainingContext(
    model=gan_model,
    optimizer=(optimizer_g, optimizer_d),  # tuple of optimizers
    fabric=fabric,
    train_metric_recorder=recorder,
    val_metric_recorder=recorder.clone() if hasattr(recorder, 'clone') else TorchMetricRecorder(None, scalar_metrics=['g_loss', 'd_loss']),
)
ctx.launch()

# Wire GAN procedure
gan_proc = BasicGANFabricProcedure(latent_dim=100)
procedure.train_loop(
    ctx, train_loader,
    max_epochs=200,
    fn_train_step=gan_proc.gan_train_step,
    fn_val_step=gan_proc.gan_val_step,
    fn_batch2target=lambda batch: batch[0],  # GAN has no labels
)
```

### Key points for GAN

- `ctx.optimizer` is a **tuple** `(optimizer_g, optimizer_d)` — the procedure unpacks it via `cast(Any, ctx.optimizer)`
- `gan_train_step` returns `(None, {'d_loss': float, 'g_loss': float})` — y_pred is None because GAN doesn't classify
- `fn_batch2target` must be overridden since default `batch[1]` doesn't apply
- GAN variants available: `StandardGAN`, `WGAN`, `WGANGP`, `LSGAN`, `HingeGAN`, `RaGAN`

## EDL: BasicEDLFabricTrainProcedure

Source: `zyplib/paradigm/edl.py:693-774`

Evidential Deep Learning has unique requirements:
- **Loss computed by the model** (`model.compute_loss(logits, y, epoch=epoch_idx)`)
- **Epoch-dependent loss** (KL regularization annealing)
- **Multiple loss components**: `total_loss`, `main_loss`, `reg_loss`, `annealing_coeff`
- **y_pred is raw logits** (not probabilities) — `before_metric` handles activation

### Setup

```python
from zyplib.paradigm.edl import EvidentialClassifier, EvidentialConfig, RegularizationConfig, BasicEDLFabricTrainProcedure
from zyplib.train import default_fabric, FabricTrainingContext, procedure, TorchMetricRecorder, clf_metrics_builder

# EDL model wraps a backbone
edl_config = EvidentialConfig(task_type='classification', num_classes=2)
reg_config = RegularizationConfig(enable_kl_regularization=True, annealing_start=10, annealing_steps=10)
edl_model = EvidentialClassifier(backbone=my_model, config=edl_config, reg_config=reg_config)

optimizer = torch.optim.Adam(edl_model.parameters(), lr=1e-3)
fabric = default_fabric(dir_name='edl_logs', keep_best_k=3)

# EDL returns multiple loss components — use 'auto' or explicit list
recorder = TorchMetricRecorder(
    clf_metrics_builder(binary=True),
    before_metric='sigmoid',
    scalar_metrics='auto',  # auto-tracks total_loss, main_loss, reg_loss, annealing_coeff
)

ctx = FabricTrainingContext(
    model=edl_model,
    optimizer=optimizer,
    fabric=fabric,
    # fn_loss is NOT set — EDL computes loss internally
    train_metric_recorder=recorder,
    val_metric_recorder=TorchMetricRecorder(
        clf_metrics_builder(binary=True),
        before_metric='sigmoid',
        scalar_metrics='auto',
    ),
)
ctx.launch()

# Wire EDL procedure
edl_proc = BasicEDLFabricTrainProcedure(edl_config)
procedure.train_loop(
    ctx, train_loader,
    max_epochs=100,
    val_loader=val_loader,
    fn_train_step=edl_proc.evidential_train_step,
    fn_val_step=edl_proc.evidential_val_step,
)
```

### Key points for EDL

- `fn_loss` is NOT set on ctx — the EDL model computes loss internally via `model.compute_loss()`
- `evidential_train_step` returns `(logits, {'total_loss': ..., 'main_loss': ..., 'reg_loss': ..., 'annealing_coeff': ...})`
- The logits are raw model output — `before_metric='sigmoid'` converts them for metric computation
- `epoch_idx` is passed through to `compute_loss` for KL annealing
- For uncertainty metrics, use `EDLUncertaintyMetric` as a custom torchmetric

## TML: TMLFabricTrainingProcedure

Source: `zyplib/paradigm/trusted_multiview.py:696-811`

Trusted Multi-view Learning has unique requirements:
- **Multi-view input**: batch is `(dict[str, Tensor], Tensor)` or `(list[Tensor], Tensor)`
- **Multiple loss components**: `total_loss`, `view_loss`, `fusion_loss`, `conflict_loss`, per-view losses
- **Model handles fusion** internally

### Setup

```python
from zyplib.paradigm.trusted_multiview import TrustedMultiViewClassifier, TrustedMultiViewConfig, TMLFabricTrainingProcedure
from zyplib.train import default_fabric, FabricTrainingContext, procedure, TorchMetricRecorder, clf_metrics_builder

config = TrustedMultiViewConfig(
    num_classes=2,
    view_names=['temporal', 'frequency'],
    input_format='dict',
    fusion_method='weighted',
)

tml_model = TrustedMultiViewClassifier(
    backbones={'temporal': temporal_net, 'frequency': freq_net},
    config=config,
)

optimizer = torch.optim.Adam(tml_model.parameters(), lr=1e-3)
fabric = default_fabric(dir_name='tml_logs', keep_best_k=3)

ctx = FabricTrainingContext(
    model=tml_model,
    optimizer=optimizer,
    fabric=fabric,
    train_metric_recorder=TorchMetricRecorder(
        clf_metrics_builder(binary=True), before_metric='sigmoid', scalar_metrics='auto',
    ),
    val_metric_recorder=TorchMetricRecorder(
        clf_metrics_builder(binary=True), before_metric='sigmoid', scalar_metrics='auto',
    ),
)
ctx.launch()

tml_proc = TMLFabricTrainingProcedure(config)
procedure.train_loop(
    ctx, train_loader,
    max_epochs=100,
    val_loader=val_loader,
    fn_train_step=tml_proc.tml_train_step,
    fn_val_step=tml_proc.tml_val_step,
)
```

### Key points for TML

- Batch format depends on `config.input_format`: `'dict'` -> `({'temporal': x1, 'frequency': x2}, y)`, `'list'` -> `([x1, x2], y)`
- `fn_batch2target` default `batch[1]` works because target is always the second element
- Returns fused probabilities as y_pred for metric computation
- `scalar_metrics='auto'` is recommended due to many dynamic loss keys

## Writing your own Procedure class

Template:

```python
class MyProcedure:
    def __init__(self, config):
        self.config = config

    def my_train_step(self, ctx, batch, *, batch_idx, epoch_idx):
        # 1. Unpack batch
        x, y = batch
        # 2. Forward
        ctx.model.train()
        ctx.optimizer.zero_grad()
        output = ctx.model(x)
        # 3. Compute loss
        loss = my_loss(output, y)
        # 4. Backward
        ctx.fabric.backward(loss)
        ctx.optimizer.step()
        # 5. Return (predictions, scalar_metrics_dict)
        return output, {'loss': loss.item()}

    def my_val_step(self, ctx, batch, *, batch_idx, epoch_idx=None):
        x, y = batch
        ctx.model.eval()
        with torch.no_grad():
            output = ctx.model(x)
            loss = my_loss(output, y)
        return output, {'loss': loss.item()}
```

Wire it:
```python
proc = MyProcedure(config)
procedure.train_loop(ctx, loader, fn_train_step=proc.my_train_step, fn_val_step=proc.my_val_step)
```
