deepdraw-opt 0.1.1__py3-none-any.whl

core/README.md

@@ -0,0 +1,14 @@
+# Core modules
+
+These modules provide Deepdraw's active-learning components. For installation and experimental use, start with the [main README](../README.md).
+
+- [experiment.py](experiment.py): retrospective active-learning loop.
+- [data_loader.py](data_loader.py): aligned measurement and embedding data for retrospective runs.
+- [initial_selection_strategies.py](initial_selection_strategies.py) and [query_strategies.py](query_strategies.py): initial and subsequent batch selection.
+- [predictor_trainer.py](predictor_trainer.py): model training and feature/target transforms.
+- [botorch_models.py](botorch_models.py) and [stable_rbf_kernel.py](stable_rbf_kernel.py): BoTorch predictors and the RBF kernel.
+- [uncertainty.py](uncertainty.py): predictive uncertainty helpers.
+- [stopping.py](stopping.py): measured-threshold and expected-improvement stopping used by both the experimental and benchmark workflows.
+- [metrics_calculator.py](metrics_calculator.py) and [round_tracker.py](round_tracker.py): retrospective round metrics and run summaries.
+
+Optional noise simulation and confirmed-top benchmark stopping are in [benchmarks/noise_simulation.py](../benchmarks/noise_simulation.py).

core/__init__.py

@@ -0,0 +1,23 @@
+"""
+Core components for active learning experiments.
+
+This package contains the refactored components that break down
+the ActiveLearningExperiment class into smaller, focused classes.
+"""
+
+from core.data_loader import DataLoader, Dataset
+from core.initial_selection_strategies import InitialSelectionStrategy
+from core.predictor_trainer import PredictorTrainer
+from core.query_strategies import QueryStrategyBase
+from core.round_tracker import RoundTracker
+from core.uncertainty import UncertaintyWrapper
+
+__all__ = [
+    "DataLoader",
+    "Dataset",
+    "PredictorTrainer",
+    "InitialSelectionStrategy",
+    "QueryStrategyBase",
+    "RoundTracker",
+    "UncertaintyWrapper",
+]

core/botorch_models.py

@@ -0,0 +1,305 @@
+"""
+BoTorch model wrappers for sklearn-style integration.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+from botorch.fit import fit_gpytorch_mll
+from botorch.models import SingleTaskGP
+from botorch.models.ensemble import EnsembleModel
+from gpytorch.constraints import GreaterThan
+from gpytorch.kernels import LinearKernel, ScaleKernel
+from gpytorch.likelihoods import GaussianLikelihood
+from gpytorch.mlls import ExactMarginalLogLikelihood
+from gpytorch.priors import GammaPrior
+from sklearn.base import RegressorMixin
+from sklearn.ensemble import RandomForestRegressor
+
+from core.stable_rbf_kernel import StableRBFKernel
+
+
+class _RandomForestEnsembleModel(EnsembleModel):
+    """
+    BoTorch ensemble wrapper around fitted sklearn tree ensembles.
+    """
+
+    def __init__(
+        self,
+        estimators: list[Any],
+        train_X: torch.Tensor,
+        weights: torch.Tensor | None = None,
+    ) -> None:
+        super().__init__(weights=weights)
+        self.estimators = list(estimators)
+        self.train_inputs = [train_X]
+        self._num_outputs = 1
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        batch_shape = X.shape[:-2]
+        q = X.shape[-2]
+        flat_X = X.detach().cpu().reshape(-1, X.shape[-1]).numpy()
+        member_values: list[torch.Tensor] = []
+        for estimator in self.estimators:
+            pred = estimator.predict(flat_X)
+            pred_tensor = torch.as_tensor(pred, device=X.device, dtype=X.dtype)
+            pred_tensor = pred_tensor.reshape(*batch_shape, q, 1)
+            member_values.append(pred_tensor)
+        return torch.stack(member_values, dim=len(batch_shape))
+
+
+class _BoTorchBaseRegressor(RegressorMixin):
+    """
+    Shared utilities for BoTorch sklearn-style wrappers.
+    """
+
+    def __init__(
+        self,
+        device: str = "cpu",
+        dtype: str = "float64",
+        ard: bool = False,
+        kernel: str = "rbf",
+    ) -> None:
+        self.device = device
+        self.dtype = dtype
+        self.ard = ard
+        self.kernel = kernel
+        self.model_: Any | None = None
+
+    def get_params(self, deep: bool = True) -> dict[str, Any]:
+        return {
+            "device": self.device,
+            "dtype": self.dtype,
+            "ard": self.ard,
+            "kernel": self.kernel,
+        }
+
+    def set_params(self, **params: Any):
+        for key, value in params.items():
+            setattr(self, key, value)
+        return self
+
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        mean, _ = self.predict_with_std(X)
+        return mean
+
+    def predict_with_std(self, X: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        self._check_is_fitted()
+        X_tensor = self._to_tensor(X)
+        with torch.no_grad():
+            posterior = self.model_.posterior(X_tensor)
+            mean = posterior.mean.squeeze(-1).cpu().numpy()
+            std = posterior.variance.clamp_min(0.0).sqrt().squeeze(-1).cpu().numpy()
+        return mean, std
+
+    def posterior(self, X: np.ndarray) -> Any:
+        self._check_is_fitted()
+        X_tensor = self._to_tensor(X)
+        return self.model_.posterior(X_tensor)
+
+    def get_botorch_model(self) -> Any:
+        self._check_is_fitted()
+        return self.model_
+
+    def _resolve_device(self) -> torch.device:
+        if self.device == "auto":
+            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        return torch.device(self.device)
+
+    def _resolve_dtype(self) -> torch.dtype:
+        if isinstance(self.dtype, torch.dtype):
+            return self.dtype
+        text = str(self.dtype).lower()
+        if text in {"float64", "double", "torch.float64"}:
+            return torch.double
+        if text in {"float32", "float", "torch.float32"}:
+            return torch.float
+        raise ValueError(f"Unsupported dtype '{self.dtype}'.")
+
+    def _build_kernel(self, X_tensor: torch.Tensor):
+        ard_dims = X_tensor.shape[-1] if self.ard else None
+        kernel_key = str(self.kernel).lower()
+        if kernel_key in {"rbf", "se", "squared_exponential"}:
+            base_kernel = StableRBFKernel(ard_num_dims=ard_dims)
+            return ScaleKernel(base_kernel)
+        elif kernel_key in {"linear"}:
+            # LinearKernel already has its own variance parameter, so wrapping it
+            # again in ScaleKernel adds a redundant output scale and makes fitting
+            # less stable on tiny high-dimensional batches.
+            return LinearKernel(ard_num_dims=ard_dims)
+        else:
+            raise ValueError(f"Unsupported kernel '{self.kernel}'. Use rbf or linear.")
+
+    def _to_tensor(self, array: np.ndarray) -> torch.Tensor:
+        device = self._resolve_device()
+        dtype = self._resolve_dtype()
+        return torch.as_tensor(array, device=device, dtype=dtype)
+
+    def _check_is_fitted(self) -> None:
+        if self.model_ is None:
+            raise ValueError(
+                f"{self.__class__.__name__}.fit must be called before predict."
+            )
+
+    def _fit_mll(self, mll: ExactMarginalLogLikelihood) -> None:
+        fit_gpytorch_mll(mll)
+
+
+class BoTorchGPRegressor(_BoTorchBaseRegressor):
+    """
+    Sklearn-compatible wrapper around a BoTorch SingleTaskGP.
+    """
+
+    def __init__(
+        self,
+        device: str = "cpu",
+        dtype: str = "float64",
+        ard: bool = False,
+        kernel: str = "rbf",
+        noise_constraint: float | None = None,
+        noise_prior: tuple[float, float] | None = None,
+    ) -> None:
+        super().__init__(
+            device=device,
+            dtype=dtype,
+            ard=ard,
+            kernel=kernel,
+        )
+        self.noise_constraint = noise_constraint
+        self.noise_prior = noise_prior
+
+    def get_params(self, deep: bool = True) -> dict[str, Any]:
+        params = super().get_params(deep=deep)
+        params.update(
+            {
+                "noise_constraint": self.noise_constraint,
+                "noise_prior": self.noise_prior,
+            }
+        )
+        return params
+
+    def fit(
+        self,
+        X: np.ndarray,
+        y: np.ndarray,
+        y_var: np.ndarray | None = None,
+    ) -> BoTorchGPRegressor:
+        X_tensor = self._to_tensor(X)
+        y_tensor = self._to_tensor(y).view(-1, 1)
+        y_var_tensor = None
+        if y_var is not None:
+            y_var_tensor = self._to_tensor(y_var).view(-1, 1)
+            y_var_tensor = y_var_tensor.clamp_min(1e-12)
+
+        covar_module = self._build_kernel(X_tensor)
+        model_kwargs = {
+            "covar_module": covar_module.to(
+                device=X_tensor.device,
+                dtype=X_tensor.dtype,
+            ),
+        }
+        if y_var_tensor is None:
+            likelihood = self._build_likelihood()
+            model_kwargs["likelihood"] = likelihood.to(
+                device=X_tensor.device,
+                dtype=X_tensor.dtype,
+            )
+        else:
+            model_kwargs["train_Yvar"] = y_var_tensor
+
+        self.model_ = SingleTaskGP(X_tensor, y_tensor, **model_kwargs)
+        self.model_.train()
+        mll = ExactMarginalLogLikelihood(self.model_.likelihood, self.model_)
+        self._fit_mll(mll)
+        self.model_.eval()
+        return self
+
+    def _build_likelihood(self) -> GaussianLikelihood:
+        constraint = (
+            GreaterThan(self.noise_constraint)
+            if self.noise_constraint is not None
+            else None
+        )
+        likelihood = (
+            GaussianLikelihood(noise_constraint=constraint)
+            if constraint is not None
+            else GaussianLikelihood()
+        )
+        if self.noise_prior is not None:
+            concentration, rate = self.noise_prior
+            likelihood.noise_covar.register_prior(
+                "noise_prior",
+                GammaPrior(concentration, rate),
+                "noise",
+            )
+        return likelihood
+
+
+class BoTorchRandomForestEnsembleRegressor(_BoTorchBaseRegressor):
+    """
+    Sklearn-compatible random forest with a BoTorch ensemble posterior.
+    """
+
+    def __init__(
+        self,
+        n_estimators: int = 100,
+        max_depth: int | None = None,
+        min_samples_split: int = 2,
+        min_samples_leaf: int = 1,
+        max_features: str | int | float | None = 1.0,
+        bootstrap: bool = True,
+        random_state: int | None = None,
+        n_jobs: int | None = None,
+        dtype: str = "float64",
+    ) -> None:
+        super().__init__(device="cpu", dtype=dtype, ard=False, kernel="rbf")
+        self.n_estimators = n_estimators
+        self.max_depth = max_depth
+        self.min_samples_split = min_samples_split
+        self.min_samples_leaf = min_samples_leaf
+        self.max_features = max_features
+        self.bootstrap = bootstrap
+        self.random_state = random_state
+        self.n_jobs = n_jobs
+        self.estimator_: RandomForestRegressor | None = None
+
+    def get_params(self, deep: bool = True) -> dict[str, Any]:
+        return {
+            "n_estimators": self.n_estimators,
+            "max_depth": self.max_depth,
+            "min_samples_split": self.min_samples_split,
+            "min_samples_leaf": self.min_samples_leaf,
+            "max_features": self.max_features,
+            "bootstrap": self.bootstrap,
+            "random_state": self.random_state,
+            "n_jobs": self.n_jobs,
+            "dtype": self.dtype,
+        }
+
+    def fit(self, X: np.ndarray, y: np.ndarray) -> BoTorchRandomForestEnsembleRegressor:
+        X_np = np.asarray(X)
+        y_np = np.asarray(y)
+        estimator = RandomForestRegressor(
+            n_estimators=self.n_estimators,
+            max_depth=self.max_depth,
+            min_samples_split=self.min_samples_split,
+            min_samples_leaf=self.min_samples_leaf,
+            max_features=self.max_features,
+            bootstrap=self.bootstrap,
+            random_state=self.random_state,
+            n_jobs=self.n_jobs,
+        )
+        estimator.fit(X_np, y_np)
+        self.estimator_ = estimator
+        train_X = torch.as_tensor(
+            X_np, device=torch.device("cpu"), dtype=self._resolve_dtype()
+        )
+        self.model_ = _RandomForestEnsembleModel(
+            estimators=list(estimator.estimators_),
+            train_X=train_X,
+        )
+        self.model_.eval()
+        return self

core/data_loader.py

@@ -0,0 +1,175 @@
+"""
+Data loading utilities for active learning experiments.
+
+This loader pairs an embeddings file with a CSV containing
+labels (and optional sample identifiers), ensuring row alignment between the two.
+"""
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Dataset:
+    """
+    Container for loaded dataset with all necessary components.
+
+    Attributes:
+        sample_ids: Stable identifiers aligned with embeddings/labels
+        labels: Array of target label values (e.g., expression)
+        embeddings: Pre-computed embeddings (required)
+        metadata: Optional metadata rows aligned to sample_ids/labels.
+    """
+
+    sample_ids: list[str]
+    labels: np.ndarray
+    embeddings: np.ndarray
+    metadata: pd.DataFrame | None = None
+
+    def __post_init__(self) -> None:
+        """Validate dataset after initialization."""
+        if len(self.sample_ids) != len(self.labels):
+            raise ValueError(
+                f"Sample IDs ({len(self.sample_ids)}) and labels "
+                f"({len(self.labels)}) must have the same length"
+            )
+        if len(self.sample_ids) != len(self.embeddings):
+            raise ValueError(
+                f"Sample IDs ({len(self.sample_ids)}) and embeddings "
+                f"({len(self.embeddings)}) must have the same length"
+            )
+
+
+class DataLoader:
+    """
+    Load embeddings from NPZ and labels from a paired CSV.
+    """
+
+    def __init__(
+        self,
+        embeddings_path: str,
+        metadata_path: str,
+        label_key: str,
+        subset_ids_path: str | None = None,
+    ) -> None:
+        """
+        Initialize the data loader.
+
+        Args:
+            embeddings_path: Path to npz file containing embeddings.
+            metadata_path: CSV with labels aligned to embeddings.
+            label_key: Column in the CSV to use as the training label.
+            subset_ids_path: Optional path to a newline-delimited file of sample IDs to keep.
+        """
+        self.embeddings_path = embeddings_path
+        self.metadata_path = metadata_path
+        self.label_key = label_key
+        self.subset_ids_path = subset_ids_path
+        self.dataset: Dataset | None = None
+
+    def load(self) -> Dataset:
+        """
+        Load paired embeddings/metadata and return a Dataset.
+        """
+        logger.info(
+            f"Loading embeddings from {self.embeddings_path} "
+            f"and metadata from {self.metadata_path}"
+        )
+
+        embeddings, sample_ids = self._load_embeddings()
+        embeddings, sample_ids = self._apply_subset_if_needed(embeddings, sample_ids)
+        labels, metadata = self._load_metadata(sample_ids)
+
+        self.dataset = Dataset(
+            sample_ids=sample_ids,
+            labels=labels,
+            embeddings=embeddings,
+            metadata=metadata,
+        )
+
+        logger.info(
+            f"Loaded dataset with {len(self.dataset.sample_ids)} samples. "
+            f"Embeddings shape: {self.dataset.embeddings.shape}"
+        )
+
+        return self.dataset
+
+    def _load_embeddings(self) -> tuple[np.ndarray, np.ndarray]:
+        data = np.load(self.embeddings_path, allow_pickle=True)
+        if "embeddings" not in data:
+            raise ValueError(
+                f"'embeddings' array not found in {self.embeddings_path}. "
+                f"Available keys: {list(data.keys())}"
+            )
+        embeddings = data["embeddings"]
+        sample_ids = data["ids"].astype(
+            np.int32, copy=False
+        )  # sample_ids is row index of csv, so we need to convert it to integer
+        return embeddings, sample_ids
+
+    def _load_metadata(self, sample_ids: np.ndarray) -> tuple[np.ndarray, pd.DataFrame]:
+        df = pd.read_csv(self.metadata_path)
+        aligned = df.iloc[sample_ids].reset_index(drop=True)
+        labels = aligned[self.label_key].to_numpy()
+        return labels, aligned
+
+    def _apply_subset_if_needed(
+        self, embeddings: np.ndarray, sample_ids: np.ndarray
+    ) -> tuple[np.ndarray, np.ndarray]:
+        subset_ids = self._load_subset_ids(sample_ids.dtype)
+        if subset_ids is None:
+            return embeddings, sample_ids
+
+        mask = np.isin(sample_ids, subset_ids)
+        if not np.any(mask):
+            raise ValueError(
+                "Subset id filtering removed all samples. "
+                "Ensure the subset ids match those stored in the embeddings file."
+            )
+
+        filtered_embeddings = embeddings[mask]
+        filtered_sample_ids = sample_ids[mask]
+
+        missing = set(subset_ids.tolist()) - set(filtered_sample_ids.tolist())
+        if missing:
+            logger.warning(
+                "Subset ids file contained %d ids not present in embeddings; ignoring smallest few: %s",
+                len(missing),
+                sorted(missing)[:5],
+            )
+
+        logger.info(
+            "Subset filtering retained %d / %d samples.",
+            len(filtered_sample_ids),
+            len(sample_ids),
+        )
+        return filtered_embeddings, filtered_sample_ids
+
+    def _load_subset_ids(self, dtype) -> np.ndarray | None:
+        if not self.subset_ids_path:
+            return None
+        subset_path = Path(self.subset_ids_path)
+        if not subset_path.exists():
+            raise FileNotFoundError(f"Subset ids file {subset_path} does not exist.")
+        subset_ids = []
+        for line in subset_path.read_text().splitlines():
+            text = line.strip()
+            if not text:
+                continue
+            try:
+                subset_ids.append(int(text))
+            except ValueError as exc:
+                raise ValueError(
+                    f"Invalid sample id '{text}' in subset file {subset_path}"
+                ) from exc
+        if not subset_ids:
+            raise ValueError(
+                f"Subset ids file {subset_path} did not contain any sample ids."
+            )
+        return np.asarray(subset_ids, dtype=dtype)