Add TabICL-v2 solver and UEA multivariate dataset loader

datasets/uea.py

@@ -0,0 +1,79 @@
+"""UEA Archive time series classification dataset.
+
+Uses tslearn to download and load any UEA multivariate dataset by name.
+Each time series is returned as a ``(T, C)`` numpy array with ``C > 1``
+channels; labels are integers.
+
+The UEA archive and the UCR archive are served by the same tslearn loader
+(``UCR_UEA_datasets``). This file is a dedicated entry point for the
+multivariate UEA datasets so they appear as a separate dataset family in
+benchopt results.
+
+Data contract output
+--------------------
+X_train : List[np.ndarray (T, C)]   one array per training sample
+y_train : np.ndarray (N,) int       class labels
+X_test  : List[np.ndarray (T, C)]
+y_test  : np.ndarray (M,) int
+task    : "classification"
+metrics : ["accuracy", "balanced_accuracy", "f1_weighted"]
+n_classes : int
+"""
+
+import numpy as np
+from benchopt import BaseDataset
+from sklearn.preprocessing import LabelEncoder
+from tslearn.datasets import UCR_UEA_datasets
+
+
+class Dataset(BaseDataset):
+    """UEA Archive multivariate classification dataset.
+
+    Parameters
+    ----------
+    dataset_name : str
+        Name of a UEA multivariate dataset (e.g. "BasicMotions",
+        "EthanolConcentration", "NATOPS").
+    debug : bool
+        If True, keep only the first 20 training samples for fast testing.
+    """
+
+    name = "UEA"
+
+    requirements = ["pip::tslearn"]
+
+    parameters = {
+        "dataset_name": ["BasicMotions"],
+        "debug": [False],
+    }
+
+    def get_data(self):
+
+        loader = UCR_UEA_datasets()
+        X_tr, y_tr, X_te, y_te = loader.load_dataset(self.dataset_name)
+
+        # tslearn returns (N, T, C) — already the right layout.
+        X_tr = np.asarray(X_tr, dtype=np.float32)
+        X_te = np.asarray(X_te, dtype=np.float32)
+
+        # Encode string labels to consecutive integers.
+        le = LabelEncoder()
+        y_tr_enc = le.fit_transform(y_tr).astype(np.int64)
+        y_te_enc = le.transform(y_te).astype(np.int64)
+
+        if self.debug:
+            X_tr = X_tr[:20]
+            y_tr_enc = y_tr_enc[:20]
+
+        # Convert to list of (T, C) arrays so variable-length datasets work too
+        X_train, X_test = list(X_tr), list(X_te)
+
+        return dict(
+            X_train=X_train,
+            y_train=y_tr_enc,
+            X_test=X_test,
+            y_test=y_te_enc,
+            task="classification",
+            metrics=["accuracy", "balanced_accuracy", "f1_weighted"],
+            n_classes=int(len(le.classes_)),
+        )

solvers/tabicl.py

@@ -0,0 +1,133 @@
+"""TabICL-v2 solver for UCR classification.
+
+TabICL is a tabular foundation model that uses in-context learning (ICL) to
+classify new samples in a single forward pass. For the current UCR scope,
+datasets are univariate and fixed-length, so each ``(T, 1)`` sample is
+reshaped into a flat ``(T,)`` feature row and the full training matrix is
+passed to ``TabICLClassifier``.
+
+References:
+    https://github.com/soda-inria/tabicl
+    https://huggingface.co/tabicl
+"""
+
+import numpy as np
+import torch
+from benchopt import BaseSolver
+
+from tabicl import TabICLClassifier
+
+SUPPORTED_TASKS = {"classification"}
+
+
+class _TabICLAdapter:
+    """Wrap a fitted TabICLClassifier behind the benchmark adapter API.
+
+    ``predict`` receives the entire test list in one call because TabICL is
+    an in-context learner — all test rows are scored jointly against the
+    stored training context.
+    """
+
+    def __init__(self, model):
+        self.model = model
+
+    def predict(self, X):
+        return self.model.predict(_to_tabular(X))
+
+
+def _to_tabular(X):
+    """Convert a list of benchmark samples into a 2-D feature matrix.
+
+    Parameters
+    ----------
+    X : list of np.ndarray, each shape (T, C)
+        UCR samples.  ``C`` is 1 for the current scope.
+
+    Returns
+    -------
+    np.ndarray, shape (N, T * C)
+        Flat tabular matrix ready for TabICL.
+    """
+    arr = np.asarray(X, dtype=np.float32)  # (N, T) or (N, T, C)
+    if arr.ndim == 3:
+        arr = arr.reshape(arr.shape[0], -1)  # flatten C into features
+    return arr
+
+
+class Solver(BaseSolver):
+    """TabICL-v2 in-context learning solver for UCR classification.
+
+    The classifier is instantiated once in ``set_objective`` (not timed) so
+    that any checkpoint download is excluded from benchmark timing. The
+    actual ``fit`` — which stores the training context — runs in ``run``.
+    """
+
+    name = "TabICL-v2"
+
+    requirements = ["pip::tabicl"]
+
+    sampling_strategy = "run_once"
+
+    parameters = {
+        "checkpoint_version": ["tabicl-classifier-v2-20260212.ckpt"],
+        "n_estimators": [8],
+    }
+
+    def skip(self, task, **kwargs):
+        if task not in SUPPORTED_TASKS:
+            return True, f"TabICL solver does not support task={task!r}"
+        return False, None
+
+    def set_objective(self, task, X_train, y_train, **meta):
+        """Prepare the solver for a given dataset configuration.
+
+        Classifier instantiation is done here (not inside ``run``) so that
+        the checkpoint download/init time is excluded from benchmark timing.
+        """
+        self.task = task
+        self.X_train = X_train
+        self.y_train = y_train
+        self.meta = meta
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Reinstantiate when the checkpoint version or n_estimators changes.
+        current_key = (self.checkpoint_version, self.n_estimators)
+        should_reload = (
+            not hasattr(self, "_classifier")
+            or not hasattr(self, "_loaded_key")
+            or self._loaded_key != current_key
+        )
+        if should_reload:
+            try:
+                self._classifier = TabICLClassifier(
+                    checkpoint_version=self.checkpoint_version,
+                    n_estimators=self.n_estimators,
+                    device=device,
+                    random_state=42,
+                    verbose=False,
+                )
+                self._loaded_key = current_key
+                print(
+                    f"\u2713 TabICL checkpoint ready: {self.checkpoint_version} "
+                    f"on device: {device}"
+                )
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to initialise TabICL checkpoint "
+                    f"'{self.checkpoint_version}': {e}. "
+                    "Make sure you have internet access and tabicl is installed."
+                ) from e
+
+        self._device = device
+        self._adapter = _TabICLAdapter(self._classifier)
+
+    def run(self, _):
+        """Fit TabICL on the training data (timed)."""
+        X_fit = _to_tabular(self.X_train)
+        y_fit = np.asarray(self.y_train)
+        self._classifier.fit(X_fit, y_fit)
+
+    def get_result(self):
+        """Return the fitted adapter wrapping the TabICL classifier."""
+        return {"model": self._adapter}

solvers/tabpfn.py

@@ -0,0 +1,127 @@
+"""TabPFN-v2 solver for UCR/UEA classification.
+
+TabPFN is a Prior-fitted Network (PFN) — a Transformer pre-trained over
+synthetic tabular datasets via meta-learning. For the UCR/UEA scope each
+``(T, C)`` time-series sample is reshaped into a flat ``(T*C,)`` feature row
+and the full training matrix is passed to ``TabPFNClassifier``.
+
+References:
+    https://github.com/PriorLabs/TabPFN
+    https://doi.org/10.1038/s41586-024-08328-6
+"""
+
+import numpy as np
+from benchopt import BaseSolver
+
+from tabpfn import TabPFNClassifier
+from tabpfn.constants import ModelVersion
+
+SUPPORTED_TASKS = {"classification"}
+
+
+class _TabPFNAdapter:
+    """Wrap a fitted TabPFNClassifier behind the benchmark adapter API.
+
+    ``predict`` receives the entire test list in one call — scoring all test
+    rows at once is significantly faster than individual calls because TabPFN
+    recomputes the training context on every call.
+    """
+
+    def __init__(self, model):
+        self.model = model
+
+    def predict(self, X):
+        return self.model.predict(_to_tabular(X))
+
+
+def _to_tabular(X):
+    """Convert a list of benchmark samples into a 2-D feature matrix.
+
+    Parameters
+    ----------
+    X : list of np.ndarray, each shape (T, C)
+        UCR/UEA samples.  ``C`` is 1 for univariate datasets.
+
+    Returns
+    -------
+    np.ndarray, shape (N, T * C)
+        Flat tabular matrix ready for TabPFN.
+    """
+    arr = np.asarray(X, dtype=np.float32)  # (N, T) or (N, T, C)
+    if arr.ndim == 3:
+        arr = arr.reshape(arr.shape[0], -1)  # flatten C into features
+    return arr
+
+
+class Solver(BaseSolver):
+    """TabPFN-v2 in-context learning solver for UCR/UEA classification.
+
+    The classifier is instantiated once in ``set_objective`` (not timed) so
+    that any checkpoint download is excluded from benchmark timing. The
+    actual ``fit`` — which stores the training context — runs in ``run``.
+    """
+
+    name = "TabPFN-v2"
+
+    requirements = ["pip::tabpfn"]
+
+    sampling_strategy = "run_once"
+
+    parameters = {
+        "n_estimators": [8],
+    }
+
+    def skip(self, task, **kwargs):
+        if task not in SUPPORTED_TASKS:
+            return True, f"TabPFN solver does not support task={task!r}"
+        return False, None
+
+    def set_objective(self, task, X_train, y_train, **meta):
+        """Prepare the solver for a given dataset configuration.
+
+        Classifier instantiation is done here (not inside ``run``) so that
+        the checkpoint download/init time is excluded from benchmark timing.
+        """
+        self.task = task
+        self.X_train = X_train
+        self.y_train = y_train
+        self.meta = meta
+
+        # Reinstantiate only when n_estimators changes.
+        should_reload = (
+            not hasattr(self, "_classifier")
+            or not hasattr(self, "_loaded_n_estimators")
+            or self._loaded_n_estimators != self.n_estimators
+        )
+        if should_reload:
+            try:
+                self._classifier = TabPFNClassifier.create_default_for_version(
+                    ModelVersion.V2,
+                    n_estimators=self.n_estimators,
+                    device="auto",
+                    random_state=42,
+                    ignore_pretraining_limits=True,
+                )
+                self._loaded_n_estimators = self.n_estimators
+                print(
+                    f"\u2713 TabPFN v2 ready "
+                    f"(n_estimators={self.n_estimators}, device=auto)"
+                )
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to initialise TabPFN v2: {e}. "
+                    "Make sure tabpfn is installed and the v2 checkpoint "
+                    "is available in the local cache (~/.cache/tabpfn/)."
+                ) from e
+
+        self._adapter = _TabPFNAdapter(self._classifier)
+
+    def run(self, _):
+        """Fit TabPFN on the training data (timed)."""
+        X_fit = _to_tabular(self.X_train)
+        y_fit = np.asarray(self.y_train)
+        self._classifier.fit(X_fit, y_fit)
+
+    def get_result(self):
+        """Return the fitted adapter wrapping the TabPFN classifier."""
+        return {"model": self._adapter}