21. Sampling API¶

@dataclass(frozen=True)
class ImbalanceSpec:
    """Optional class imbalance scenario.

    Kinds:
    - none
    - subsample_max_per_class: cap each class to max_per_class (applies to train or labeled)
    - long_tail: exponential decay per class rank (applies to train or labeled)

    apply_to:
    - train: modify train_idx before labeling
    - labeled: modify labeled subset after labeling (removed labeled become unlabeled)
    """

    kind: Literal["none", "subsample_max_per_class", "long_tail"] = "none"
    apply_to: Literal["train", "labeled"] = "train"
    max_per_class: int | None = None
    alpha: float | None = None
    min_per_class: int = 1

    def as_dict(self) -> dict[str, Any]:
        return {
            "kind": self.kind,
            "apply_to": self.apply_to,
            "max_per_class": None if self.max_per_class is None else int(self.max_per_class),
            "alpha": None if self.alpha is None else float(self.alpha),
            "min_per_class": int(self.min_per_class),
        }

    @classmethod
    def from_dict(cls, d: Mapping[str, Any]) -> ImbalanceSpec:
        _assert_known_keys(
            d,
            {"kind", "apply_to", "max_per_class", "alpha", "min_per_class"},
            "imbalance",
        )
        kind = str(d.get("kind", "none"))
        if kind not in ("none", "subsample_max_per_class", "long_tail"):
            raise ValueError(f"Unknown imbalance kind: {kind!r}")
        apply_to = str(d.get("apply_to", "train"))
        if apply_to not in ("train", "labeled"):
            raise ValueError(f"Unknown imbalance apply_to: {apply_to!r}")
        return cls(
            kind=kind,  # type: ignore[arg-type]
            apply_to=apply_to,  # type: ignore[arg-type]
            max_per_class=d.get("max_per_class", None),
            alpha=d.get("alpha", None),
            min_per_class=int(d.get("min_per_class", 1)),
        )

@dataclass(frozen=True)
class LabelingSpec:
    """How to select labeled samples within the train partition.

    Modes:
    - fraction: value in (0, 1], selects that fraction of train samples
    - count: value is an integer count of labeled samples
    - per_class: value is an integer count per class

    If fixed_indices is provided, it is used directly (validated) and the mode is ignored.
    """

    mode: Literal["fraction", "count", "per_class"] = "fraction"
    value: float | int = 0.1
    per_class: bool = False
    min_per_class: int = 1
    strategy: Literal["proportional", "balanced"] = "proportional"
    fixed_indices: Sequence[int] | None = None

    def as_dict(self) -> dict[str, Any]:
        return {
            "mode": self.mode,
            "value": float(self.value) if self.mode == "fraction" else int(self.value),
            "per_class": bool(self.per_class),
            "min_per_class": int(self.min_per_class),
            "strategy": self.strategy,
            "fixed_indices": None
            if self.fixed_indices is None
            else [int(i) for i in self.fixed_indices],
        }

    @classmethod
    def from_dict(cls, d: Mapping[str, Any]) -> LabelingSpec:
        _assert_known_keys(
            d,
            {"mode", "value", "per_class", "min_per_class", "strategy", "fixed_indices"},
            "labeling",
        )
        mode = str(d.get("mode", "fraction"))
        if mode not in ("fraction", "count", "per_class"):
            raise ValueError(f"Unknown labeling mode: {mode!r}")
        value = d.get("value", 0.1)
        value = float(value) if mode == "fraction" else int(value)
        fixed_indices = d.get("fixed_indices", None)
        if fixed_indices is not None:
            if isinstance(fixed_indices, (str, bytes)) or not isinstance(fixed_indices, Sequence):
                raise ValueError("labeling.fixed_indices must be a sequence of integers")
            fixed_indices = [int(i) for i in fixed_indices]
        strategy = str(d.get("strategy", "proportional"))
        if strategy not in ("proportional", "balanced"):
            raise ValueError(f"Unknown labeling strategy: {strategy!r}")
        return cls(
            mode=mode,  # type: ignore[arg-type]
            value=value,
            per_class=bool(d.get("per_class", False)),
            min_per_class=int(d.get("min_per_class", 1)),
            strategy=strategy,  # type: ignore[arg-type]
            fixed_indices=fixed_indices,
        )

class SamplingError(RuntimeError):
    """Base error for sampling."""

@dataclass(frozen=True)
class SamplingPlan:
    """Full sampling plan."""

    split: SplitSpec = field(default_factory=HoldoutSplitSpec)
    labeling: LabelingSpec = field(default_factory=LabelingSpec)
    imbalance: ImbalanceSpec = field(default_factory=ImbalanceSpec)
    policy: SamplingPolicy = field(default_factory=SamplingPolicy)

    def as_dict(self) -> dict[str, Any]:
        return {
            "split": self.split.as_dict(),
            "labeling": self.labeling.as_dict(),
            "imbalance": self.imbalance.as_dict(),
            "policy": {
                "respect_official_test": bool(self.policy.respect_official_test),
                "use_official_graph_masks": bool(self.policy.use_official_graph_masks),
                "allow_override_official": bool(self.policy.allow_override_official),
            },
        }

    @classmethod
    def from_dict(cls, d: Mapping[str, Any]) -> SamplingPlan:
        _assert_known_keys(d, {"split", "labeling", "imbalance", "policy"}, "plan")
        split_obj = _ensure_mapping(d.get("split", {}), "split")
        split_kind = str(split_obj.get("kind", "holdout"))
        if split_kind == "kfold":
            split = KFoldSplitSpec.from_dict(split_obj)
        elif split_kind == "holdout":
            split = HoldoutSplitSpec.from_dict(split_obj)
        else:
            raise ValueError(f"Unknown split kind: {split_kind!r}")

        labeling_obj = _ensure_mapping(d.get("labeling", {}), "labeling")
        labeling = LabelingSpec.from_dict(labeling_obj)

        imbalance_obj = _ensure_mapping(d.get("imbalance", {}), "imbalance")
        imbalance = ImbalanceSpec.from_dict(imbalance_obj)

        policy_obj = _ensure_mapping(d.get("policy", {}), "policy")
        policy = SamplingPolicy.from_dict(policy_obj)

        return cls(split=split, labeling=labeling, imbalance=imbalance, policy=policy)

@dataclass(frozen=True)
class SamplingPolicy:
    """Policy for handling official provider splits.

    - respect_official_test: if dataset.test exists, keep it as the test set
    - use_official_graph_masks: if graph dataset provides masks, use them as train/val/test masks
    """

    respect_official_test: bool = True
    use_official_graph_masks: bool = True
    allow_override_official: bool = False

    @classmethod
    def from_dict(cls, d: Mapping[str, Any]) -> SamplingPolicy:
        _assert_known_keys(
            d,
            {"respect_official_test", "use_official_graph_masks", "allow_override_official"},
            "policy",
        )
        return cls(
            respect_official_test=bool(d.get("respect_official_test", True)),
            use_official_graph_masks=bool(d.get("use_official_graph_masks", True)),
            allow_override_official=bool(d.get("allow_override_official", False)),
        )

@dataclass(frozen=True)
class SamplingResult:
    """Sampling result with indices (inductive) or masks (graph transductive).

    Indices keys (typical):
    - train, val, test
    - train_labeled, train_unlabeled

    Refs indicate the base split each index array refers to:
    - "train" means indices are relative to dataset.train
    - "test" means indices are relative to dataset.test
    - "nodes" means graph nodes

    Masks keys (graph):
    - train, val, test, labeled, unlabeled
    """

    schema_version: int
    created_at: str
    dataset_fingerprint: str
    split_fingerprint: str
    plan: Mapping[str, Any]

    indices: Mapping[str, np.ndarray] = field(default_factory=dict)
    refs: Mapping[str, str] = field(default_factory=dict)
    masks: Mapping[str, np.ndarray] = field(default_factory=dict)

    stats: Mapping[str, Any] = field(default_factory=dict)

    def is_graph(self) -> bool:
        return bool(self.masks)

    def validate(self, *, n_train: int, n_test: int | None, n_nodes: int | None) -> None:
        if self.is_graph():
            self._validate_graph(n_nodes=n_nodes)
        else:
            self._validate_inductive(n_train=n_train, n_test=n_test)

    def _validate_inductive(self, *, n_train: int, n_test: int | None) -> None:
        def _check_idx(name: str, base: str, idx: np.ndarray) -> None:
            if idx.dtype.kind not in ("i", "u"):
                raise SamplingValidationError(f"{name} indices must be integers, got {idx.dtype}")
            if idx.size == 0:
                return
            max_ok = n_train if base == "train" else n_test
            if max_ok is None:
                raise SamplingValidationError("n_test is required when base='test'")
            if idx.min() < 0 or idx.max() >= max_ok:
                raise SamplingValidationError(f"{name} has out-of-range indices for base={base!r}")
            if np.unique(idx).size != idx.size:
                raise SamplingValidationError(f"{name} contains duplicate indices")

        for name, idx in self.indices.items():
            base = self.refs.get(name, "train")
            _check_idx(name, base, idx)

        # core invariants
        train = self.indices.get("train")
        val = self.indices.get("val")
        test = self.indices.get("test")
        if train is None:
            raise SamplingValidationError("Missing train indices")
        if val is None:
            raise SamplingValidationError("Missing val indices")
        if test is None:
            raise SamplingValidationError("Missing test indices (may be empty array)")

        if self.refs.get("train", "train") != "train" or self.refs.get("val", "train") != "train":
            raise SamplingValidationError("train and val indices must be relative to dataset.train")

        # disjointness of train/val in same base
        if np.intersect1d(train, val).size:
            raise SamplingValidationError("train and val overlap")

        # test disjointness only if same base
        if self.refs.get("test", "train") == "train" and (
            np.intersect1d(train, test).size or np.intersect1d(val, test).size
        ):
            raise SamplingValidationError("test overlaps with train or val")

        labeled = self.indices.get("train_labeled")
        unlabeled = self.indices.get("train_unlabeled")
        if labeled is None or unlabeled is None:
            raise SamplingValidationError("Missing labeled/unlabeled indices")
        if np.intersect1d(labeled, unlabeled).size:
            raise SamplingValidationError("labeled and unlabeled overlap")
        if not np.array_equal(np.sort(np.concatenate([labeled, unlabeled])), np.sort(train)):
            raise SamplingValidationError("labeled + unlabeled must cover train exactly")

    def _validate_graph(self, *, n_nodes: int | None) -> None:
        if n_nodes is None:
            raise SamplingValidationError("n_nodes is required to validate graph masks")
        required = {"train", "val", "test", "labeled", "unlabeled"}
        if set(self.masks.keys()) != required:
            raise SamplingValidationError(f"Graph masks must have keys {sorted(required)}")

        for k, m in self.masks.items():
            if m.dtype != bool:
                raise SamplingValidationError(f"Mask {k!r} must be bool, got {m.dtype}")
            if m.shape != (n_nodes,):
                raise SamplingValidationError(
                    f"Mask {k!r} must have shape ({n_nodes},), got {m.shape}"
                )

        train = self.masks["train"]
        val = self.masks["val"]
        test = self.masks["test"]
        labeled = self.masks["labeled"]
        unlabeled = self.masks["unlabeled"]

        if (labeled & ~train).any():
            raise SamplingValidationError("labeled_mask must be subset of train_mask")
        if not np.array_equal(unlabeled, train & ~labeled):
            raise SamplingValidationError("unlabeled_mask must equal train_mask & ~labeled_mask")

        # train/val/test should be disjoint
        if (train & val).any() or (train & test).any() or (val & test).any():
            raise SamplingValidationError("train/val/test masks must be disjoint")

    @property
    def train_idx(self) -> np.ndarray:
        if self.is_graph():
            return np.where(self.masks["train"])[0]
        return self.indices.get("train", np.array([]))

    @property
    def val_idx(self) -> np.ndarray:
        if self.is_graph():
            return np.where(self.masks["val"])[0]
        return self.indices.get("val", np.array([]))

    @property
    def test_idx(self) -> np.ndarray:
        if self.is_graph():
            return np.where(self.masks["test"])[0]
        return self.indices.get("test", np.array([]))

    @property
    def labeled_idx(self) -> np.ndarray:
        if self.is_graph():
            return np.where(self.masks["labeled"])[0]
        return self.indices.get("train_labeled", np.array([]))

    @property
    def unlabeled_idx(self) -> np.ndarray:
        if self.is_graph():
            return np.where(self.masks["unlabeled"])[0]
        return self.indices.get("train_unlabeled", np.array([]))

class SamplingValidationError(SamplingError):
    """Raised when a sampled split violates invariants."""