# Build your own classifier

ARCADE ranks discriminators under one shared pipeline. Adding a method means
writing a small Python class that implements the classifier contract,
registering it under a short name, and listing that name in YAML. Splits,
optional duration sweep, analytical hardware estimation, and the summary
report stay the same as for built-in methods.

This guide covers three patterns:

1. A minimal classical toy on filter features.
2. How to leverage custom classes from YAML and from Python imports.
3. A transformer-style custom class that consumes time-resolved sequences.

A checked-in toy lives under ``examples/custom_classifier/``.

```{mermaid}
flowchart TB
  subgraph Your code
    CLS["@register_classifier name<br/>BaseClassifier subclass"]
  end
  subgraph ARCADE pipeline
    DATA[Data stage] --> FEAT["Feature build<br/>FEATURE_KIND"]
    FEAT --> FIT[fit / predict]
    FIT --> REP[Report row]
  end
  CLS --> FEAT
  CLS --> FIT
```

## Contract you must satisfy

| Piece | Role |
|-------|------|
| ``@register_classifier("name")`` | Puts the class in the registry for YAML ``classifiers.run`` |
| ``FEATURE_KIND`` | Tells the classify stage which matrix to build |
| ``fit`` / ``predict`` | Train and infer |
| ``from_config`` | Optional; maps a YAML option block into constructor arguments |
| ``CATEGORY`` / ``DESCRIPTION`` | Shown by ``arcade list`` and ``arcade describe`` |

Supported feature kinds you will use most often are ``filter_features``,
``iq_trace``, and ``raw_trace_seq``. Wrong kinds usually surface as shape errors
rather than clear messages, so set ``FEATURE_KIND`` deliberately.

## Pattern A: nearest-mean on filter features

```python
# my_pkg/nearest_mean.py
from __future__ import annotations

import numpy as np
from arcade.classifiers.base import BaseClassifier, register_classifier


@register_classifier("nearest_mean")
class NearestMeanClassifier(BaseClassifier):
    DESCRIPTION = "Nearest class-mean on filter features"
    CATEGORY = "classical"
    FEATURE_KIND = "filter_features"

    def __init__(self) -> None:
        self._means: dict[int, np.ndarray] = {}

    def fit(self, train_features, train_labels, **kwargs):
        X = np.asarray(train_features)
        y = np.asarray(train_labels).astype(int)
        self._means = {int(c): X[y == c].mean(axis=0) for c in np.unique(y)}
        return self

    def predict(self, features):
        X = np.asarray(features)
        labels = sorted(self._means)
        means = np.stack([self._means[c] for c in labels], axis=0)
        d = ((X[:, None, :] - means[None, :, :]) ** 2).sum(axis=-1)
        return np.asarray(labels, dtype=int)[d.argmin(axis=1)]

    @classmethod
    def from_config(cls, config: dict) -> "NearestMeanClassifier":
        return cls()
```

Ensure the module is importable before ``arcade run``. Install your package into
the same environment, or set ``PYTHONPATH`` so Python can find ``my_pkg``.

### Leverage the class from YAML

Registered name after the import side effect:

```yaml
classifiers:
  run:
    - threshold
    - lda
    - nearest_mean
  filter_types: [MF, RMF]
```

Dotted path when you do not want to rely on prior registration. The loader
imports the class from the path string:

```yaml
classifiers:
  run:
    - threshold
    - my_pkg.nearest_mean.NearestMeanClassifier
```

### Leverage the class from Python

```python
from arcade.classifiers.base import get_classifier
import my_pkg.nearest_mean  # registration side effect

clf = get_classifier("nearest_mean").from_config({})
# or pass the class into a notebook cell that builds features yourself
```

You can also keep the class unregistered and pass the dotted path only in YAML.
That is useful when you are prototyping inside a private package and do not want
a global short name yet.

## Pattern B: transformer-style custom class

The built-in ``transformer`` method already consumes ``raw_trace_seq`` features
and trains a ``TraceTransformer`` through ``arcade.nn.builder.build_nn``. Use it
when the stock path is enough:

```yaml
classifiers:
  run: [threshold, transformer]
  transformer:
    input_mode: raw
    nn_config:
      model:
        type: transformer
        d_model: 64
        nhead: 4
        num_layers: 2
        dim_feedforward: 128
        dropout: 0.1
      training:
        epochs: 50
        batch_size: 512
        lr: 1.0e-3
        device: auto
```

When you need a custom variant, for example a different default input mode, an
extra preprocessing step, or a renamed experimental method, subclass the same
contract and keep ``FEATURE_KIND = "raw_trace_seq"`` so the pipeline builds
time-resolved sequences.

```python
# my_pkg/tiny_transformer.py
from __future__ import annotations

from typing import Any

import numpy as np
from arcade.classifiers.base import BaseClassifier, register_classifier


@register_classifier("tiny_transformer")
class TinyTransformerClassifier(BaseClassifier):
    """Minimal attention classifier on raw IQ sequences."""

    DESCRIPTION = "Custom tiny transformer on raw_trace_seq features"
    CATEGORY = "neural"
    FEATURE_KIND = "raw_trace_seq"
    REQUIRES_NN_CONFIG = True

    def __init__(self, *, nn_config: dict | None = None) -> None:
        self.nn_config = nn_config or {}
        self._model = None
        self._seq_len = 0
        self._feature_dim = 0

    def fit(
        self,
        train_features: np.ndarray,
        train_labels: np.ndarray,
        val_features: np.ndarray | None = None,
        val_labels: np.ndarray | None = None,
        **kwargs: Any,
    ) -> "TinyTransformerClassifier":
        from arcade.nn.builder import build_nn
        from arcade.training.transformer_training import train_trace_transformer

        train_seq = np.asarray(train_features, dtype=np.float32)
        if train_seq.ndim != 3:
            raise ValueError(f"expected (n,T,F), got {train_seq.shape}")
        self._seq_len = int(train_seq.shape[1])
        self._feature_dim = int(train_seq.shape[2])
        num_classes = int(np.asarray(train_labels).max()) + 1

        cfg = dict(self.nn_config)
        model_cfg = dict(cfg.get("model", cfg))
        train_cfg = dict(cfg.get("training", {}))
        model_cfg.setdefault("type", "transformer")

        self._model = build_nn(
            {"model": model_cfg},
            seq_len=self._seq_len,
            feature_dim=self._feature_dim,
            output_dim=num_classes,
        )
        result = train_trace_transformer(
            self._model,
            train_seq,
            np.asarray(train_labels).ravel(),
            np.asarray(val_features, dtype=np.float32) if val_features is not None else None,
            np.asarray(val_labels).ravel() if val_labels is not None else None,
            epochs=int(train_cfg.get("epochs", 30)),
            batch_size=int(train_cfg.get("batch_size", 256)),
            lr=float(train_cfg.get("lr", 1e-3)),
            device=str(train_cfg.get("device", "auto")),
        )
        self._model = result["model"]
        return self

    def predict(self, features: np.ndarray) -> np.ndarray:
        import torch

        if self._model is None:
            raise RuntimeError("Call fit() before predict()")
        seq = np.asarray(features, dtype=np.float32)
        self._model.eval()
        dev = next(self._model.parameters()).device
        x = torch.tensor(seq, dtype=torch.float32, device=dev)
        with torch.no_grad():
            return self._model(x).argmax(dim=-1).cpu().numpy().astype(np.int64)

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "TinyTransformerClassifier":
        nn_cfg = config.get("nn_config")
        if nn_cfg is None and config.get("config_file"):
            import yaml

            with open(config["config_file"]) as f:
                nn_cfg = yaml.safe_load(f)
        return cls(nn_config=nn_cfg)
```

```{mermaid}
flowchart LR
  RAW[Prepared traces] --> SEQ["raw_trace_seq<br/>(n shots × T × F)"]
  SEQ --> TT[tiny_transformer fit]
  TT --> MET[Metrics + report row]
  TT --> HW[Analytical hardware]
```

### YAML that runs the custom transformer beside a baseline

```yaml
classifiers:
  run:
    - threshold
    - tiny_transformer
  tiny_transformer:
    nn_config:
      model:
        type: transformer
        d_model: 64
        nhead: 4
        num_layers: 2
        dim_feedforward: 128
      training:
        epochs: 30
        batch_size: 256
        lr: 1.0e-3
        device: auto

visualization:
  output_dir: output/custom_transformer_bakeoff
```

Install ``torch`` first:

```bash
pip install -e ".[torch]"
export PYTHONPATH=path/to/my_pkg_parent:$PYTHONPATH
arcade run path/to/config.yaml
```

The custom method appears in the same summary report as ``threshold``. That is
the point of the registry: your experiment gets the shared recipe for free.

## What you get for free

The same train, validation, and test cuts as every other method in the run.
Feature matrices selected by ``FEATURE_KIND``. Optional duration sweep when
enabled. Analytical FPGA estimates when hardware estimation is enabled. A row
in the summary report next to zoo methods.

## Pitfalls

<div class="arcade-pitfall">
<ul>
<li>Match <code>FEATURE_KIND</code> to what <code>fit</code> expects. Sequence models need
  <code>raw_trace_seq</code>; filter models need <code>filter_features</code>.</li>
<li>Neural methods need the <code>torch</code> extra and usually an <code>nn_config</code> or
  <code>config_file</code>.</li>
<li>Import the module before <code>arcade run</code>, or use a dotted class path in YAML.</li>
<li>Do not mix a custom optimized net and an unoptimized zoo net in one ranking
  table without labeling the rows.</li>
</ul>
</div>

For registry names and method-level pipeline notes, see
{doc}`../reference/classifiers`.