# `pytorch_wrapper`

Convert a trained PyTorch model into an evaluator compatible with URANIE Relauncher.

```python
uranie_evaluator(model_torch, outputkey=[])
```

**Parameters:**

```{list-table}
:header-rows: 1
:widths: 18 20 62

* - Parameter
  - Type
  - Description
* - `model_torch`
  - `torch.nn.Module`
  - A trained PyTorch model. The model must be trained (validated by checking for non-None gradients on parameters).
* - `outputkey`
  - `list`
  - Optional list of keys or indices to select specific outputs from multi-output models. For tuple/list outputs, use integers (e.g., `[0, 1]`). For dict outputs, use strings (e.g., `["mean", "std"]`). Default: `[]` (returns all outputs).
```

**Returns:**

| Type | Description |
|------|-------------|
| `callable` or `None` | A Python callable that accepts `*args` (floats) and returns a 1-D list of floats. Returns `None` if the model is not trained or arguments are invalid. |

**Exceptions/Errors:**

- Returns `None` and prints error message if `model_torch` is not a valid `torch.nn.Module`
- Returns `None` and prints error message if model has not been trained (no gradients)
- Returns `None` and prints error message if `outputkey` is invalid
- Prints warning if multi-dimensional outputs are flattened

## Examples

**One-dimensional example:**

```python
import torch
import numpy as np

from uratools.pytorch_wrapper import uranie_evaluator

# Define model
class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(4, 8)
        self.fc2 = torch.nn.Linear(8, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# Create model and optimizer
model = Net()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.MSELoss()

# Generate dummy training data
X_train = torch.randn(50, 4)
y_train = torch.randn(50, 1)

# Training loop
model.train()
for epoch in range(10):
    optimizer.zero_grad()
    predictions = model(X_train)
    loss = criterion(predictions, y_train)
    loss.backward()
    optimizer.step()

# Wrap for URANIE
wrapped = uranie_evaluator(model)
if wrapped:
    result = wrapped(0.1, 0.2, 0.3, 0.4)
    print(result)  # [predicted_value]
```

**Multi-dimensional example:**

```python
import torch
from uratools.pytorch_wrapper import uranie_evaluator

class MultiOut(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(4, 16)
        self.out1 = torch.nn.Linear(16, 2)
        self.out2 = torch.nn.Linear(16, 1)
    
    def forward(self, x):
        h = torch.relu(self.fc(x))
        return (self.out1(h), self.out2(h))

# Create model and train
model = MultiOut()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# Generate dummy training data
X_train = torch.randn(50, 4)
y1_train = torch.randn(50, 2)
y2_train = torch.randn(50, 1)

# Training loop
model.train()
for epoch in range(10):
    optimizer.zero_grad()
    out1, out2 = model(X_train)
    loss1 = torch.nn.functional.mse_loss(out1, y1_train)
    loss2 = torch.nn.functional.mse_loss(out2, y2_train)
    loss = loss1 + loss2
    loss.backward()
    optimizer.step()

# Select only first output
wrapped = uranie_evaluator(model, outputkey=[0])
if wrapped:
    result = wrapped(0.1, 0.2, 0.3, 0.4)
    print(result)  # [predicted_value_1, predicted_value_2]
```

## Implementation Details

- **Input conversion**: Arguments are converted to a 1-row float32 tensor before passing to the model.
- **Output handling**: Outputs are converted to Python lists. Multi-dimensional outputs are flattened.
- **Training check**: Model is considered trained if at least one parameter has non-None gradients.
- **Inference mode**: Model is called in `torch.no_grad()` context (inference mode).

## Supported Model Types

Works with any `torch.nn.Module` subclass that has been trained:

- `torch.nn.Linear`, `torch.nn.Sequential`, `torch.nn.ModuleList`
- `torch.nn.Conv1d`, `torch.nn.Conv2d`, `torch.nn.Conv3d`
- `torch.nn.RNN`, `torch.nn.LSTM`, `torch.nn.GRU`
- `torch.nn.TransformerEncoder`, `torch.nn.TransformerDecoder`
- Custom `torch.nn.Module` subclasses
- Pre-trained models from `torchvision`, `torchtext`, `torchaudio`
- Multi-output models (with `outputkey` parameter)

## Supported Output Types

- Scalar outputs (int, float)
- 1-D arrays, tuples, and dictionaries
- Multi-dimensional tensors
- Using `outputkey` parameter for multi-output models