augmask 0.1.0__py3-none-any.whl

augmask/__init__.py

@@ -0,0 +1,5 @@
+"""AugMask: training diffusion models on incomplete tabular data (ICML 2026)."""
+from .augment import Augmenter, Encoded, LGBMSAugmenter, MeanModeFill
+
+__version__ = "0.1.0"
+__all__ = ["Augmenter", "Encoded", "LGBMSAugmenter", "MeanModeFill"]

augmask/augment.py

@@ -0,0 +1,192 @@
+"""Step 5: fill the missing entries of the (encoded) inputs.
+
+``LGBMSAugmenter`` is the paper's factorized conditional stochastic augmentation
+(Algorithm 1, "LGBM-S"); ``MeanModeFill`` is the mean/mode fill of NoAug and MeanMask.
+Any object with the ``Augmenter`` interface can be used instead.
+Both are fit on the train split and fill train and any further splits (val).
+"""
+from __future__ import annotations
+
+import os
+from concurrent.futures import ThreadPoolExecutor
+from typing import NamedTuple, Protocol, Sequence
+
+import numpy as np
+import pandas as pd
+from lightgbm import LGBMClassifier, LGBMRegressor
+
+# -E[log chi^2_1]: for Gaussian residuals r, E[log r^2] = log sigma^2 - LOG_CHI2_BIAS
+LOG_CHI2_BIAS = float(np.euler_gamma + np.log(2))
+
+# Table 4 of the paper; the last four entries only set threading, determinism and logging.
+LGBM_PARAMS = dict(
+    boosting_type="gbdt",
+    num_leaves=31,
+    max_depth=-1,
+    learning_rate=0.1,
+    n_estimators=50,
+    subsample_for_bin=200000,
+    min_split_gain=0.0,
+    min_child_weight=1e-3,
+    min_child_samples=20,
+    subsample=1.0,
+    subsample_freq=0,
+    colsample_bytree=1.0,
+    reg_alpha=0.0,
+    reg_lambda=0.0,
+    random_state=42,
+    importance_type="split",
+    n_jobs=1,
+    deterministic=True,
+    force_row_wise=True,
+    verbosity=-1,
+)
+
+
+class Encoded(NamedTuple):
+    """One split in encoded space. Masks are True where observed."""
+
+    X_num: np.ndarray  # float, quantile-normal values (NaN where missing)
+    X_cat: np.ndarray  # int codes 0..K_j-1 (-1 where missing)
+    M_num: np.ndarray  # bool
+    M_cat: np.ndarray  # bool
+
+
+class Augmenter(Protocol):
+    """Fills the missing entries: fit on ``train``, then return filled (X_num, X_cat) for ``train``
+    followed by each of ``others``. Observed entries stay unchanged; a missing numerical entry gets a
+    float and a missing categorical entry a code in 0..K_j-1. With several fills per scenario
+    (``n_fills`` > 1), ``augment`` is called once per fill and should return a new random draw each time."""
+
+    def augment(self, train: Encoded, others: Sequence[Encoded], n_categories: Sequence[int]
+                ) -> list[tuple[np.ndarray, np.ndarray]]: ...
+
+
+class MeanModeFill:
+    """Train mean (numerical) and train mode (categorical) of the observed entries."""
+
+    def augment(self, train: Encoded, others: Sequence[Encoded], n_categories: Sequence[int]
+                ) -> list[tuple[np.ndarray, np.ndarray]]:
+        means = np.nan_to_num(np.nanmean(np.where(train.M_num, train.X_num, np.nan), axis=0))
+        modes = np.array([np.bincount(train.X_cat[train.M_cat[:, k], k], minlength=max(K, 1)).argmax()
+                          for k, K in enumerate(n_categories)], dtype=np.int64)
+        return [(np.where(b.M_num, b.X_num, means), np.where(b.M_cat, b.X_cat, modes)) for b in (train, *others)]
+
+
+def _design(block: Encoded, n_categories: Sequence[int]) -> pd.DataFrame:
+    """LightGBM inputs u = (x~, m): values with NaN where missing, plus mask indicators."""
+    cols = {}
+    for k in range(block.X_num.shape[1]):
+        cols[f"n{k}"] = np.where(block.M_num[:, k], block.X_num[:, k], np.nan)
+    for k, K in enumerate(n_categories):
+        codes = np.where(block.M_cat[:, k], block.X_cat[:, k], -1)
+        cols[f"c{k}"] = pd.Categorical.from_codes(codes, categories=np.arange(K))
+    for k in range(block.M_num.shape[1]):
+        cols[f"mn{k}"] = block.M_num[:, k].astype(np.uint8)
+    for k in range(block.M_cat.shape[1]):
+        cols[f"mc{k}"] = block.M_cat[:, k].astype(np.uint8)
+    return pd.DataFrame(cols)
+
+
+class LGBMSAugmenter:
+    """Factorized conditional stochastic augmentation with per-feature LightGBM models.
+
+    For each feature j, models are fit on the train rows where j is observed, with
+    inputs u = (x~_{-j}, m_{-j}). Numerical features are sampled from
+    N(mu_j(u), exp(f_j(u) + c)), where f_j regresses log(r^2 + eps) on u and r are the
+    residuals of mu_j. The default c = 0 is Algorithm 1; ``log_var_offset=LOG_CHI2_BIAS``
+    removes the bias of the log-scale fit (Harvey, 1976), which widens the fills by about
+    1.9x without consistently improving the generators. Categorical features are sampled from the predicted class
+    probabilities after label smoothing with alpha over the K_j categories.
+    Every model conditions on the incomplete data as observed (no chaining).
+    Features are fit in parallel on ``n_jobs`` threads (-1: all cores); the fills do not depend on it.
+    ``draw`` indexes independent fills from the same models (draw 0 is the default stream).
+    """
+
+    def __init__(self, alpha: float = 0.05, eps: float = 1e-6, log_var_offset: float = 0.0,
+                 seed: int = 42, n_jobs: int = -1, lgbm_params: dict | None = None, draw: int = 0):
+        self.alpha = alpha
+        self.eps = eps
+        self.log_var_offset = log_var_offset
+        self.seed = seed
+        self.draw = draw
+        self.n_jobs = (os.cpu_count() or 1) if n_jobs == -1 else n_jobs
+        self.lgbm_params = {**LGBM_PARAMS, **(lgbm_params or {})}
+
+    def augment(self, train: Encoded, others: Sequence[Encoded], n_categories: Sequence[int]
+                ) -> list[tuple[np.ndarray, np.ndarray]]:
+        """Return filled (X_num, X_cat) for ``train`` followed by each of ``others``."""
+        blocks = [train, *others]
+        designs = [_design(b, n_categories) for b in blocks]
+        d_num = train.X_num.shape[1]
+        features = [("num", k) for k in range(d_num)] + [("cat", k) for k in range(len(n_categories))]
+        root = np.random.SeedSequence(self.seed if self.draw == 0 else [self.seed, self.draw])
+        rngs = [np.random.default_rng(s) for s in root.spawn(len(features))]
+
+        def run(i):
+            kind, k = features[i]
+            return self._fill_feature(kind, k, blocks, designs, n_categories, rngs[i])
+
+        if self.n_jobs > 1:
+            with ThreadPoolExecutor(self.n_jobs) as pool:
+                results = list(pool.map(run, range(len(features))))
+        else:
+            results = [run(i) for i in range(len(features))]
+
+        out = []
+        for b_idx, block in enumerate(blocks):
+            X_num = np.where(block.M_num, block.X_num, np.nan)
+            X_cat = np.where(block.M_cat, block.X_cat, -1)
+            for (kind, k), fills in zip(features, results):
+                if fills is None:
+                    continue
+                if kind == "num":
+                    X_num[~block.M_num[:, k], k] = fills[b_idx]
+                else:
+                    X_cat[~block.M_cat[:, k], k] = fills[b_idx]
+            out.append((X_num, X_cat))
+        return out
+
+    def _fill_feature(self, kind, k, blocks, designs, n_categories, rng):
+        num = kind == "num"
+        name = f"n{k}" if num else f"c{k}"
+        drop = [name, f"m{name}"]
+        masks = [b.M_num[:, k] if num else b.M_cat[:, k] for b in blocks]
+        if not any((~m).any() for m in masks):
+            return None
+        obs = masks[0]
+        X_fit = designs[0].loc[obs].drop(columns=drop)
+        X_miss = [d.loc[~m].drop(columns=drop) for d, m in zip(designs, masks)]
+
+        if num:
+            y = blocks[0].X_num[obs, k]
+            mu_model = LGBMRegressor(**self.lgbm_params).fit(X_fit, y)
+            log_var = np.log((y - mu_model.predict(X_fit)) ** 2 + self.eps)
+            var_model = LGBMRegressor(**self.lgbm_params).fit(X_fit, log_var)
+            fills = []
+            for X in X_miss:
+                if len(X) == 0:
+                    fills.append(np.empty(0))
+                    continue
+                mu, sigma = mu_model.predict(X), np.exp(0.5 * (var_model.predict(X) + self.log_var_offset))
+                fills.append(mu + sigma * rng.standard_normal(len(X)))
+            return fills
+
+        K = n_categories[k]
+        y = blocks[0].X_cat[obs, k]
+        classes = np.unique(y)
+        clf = LGBMClassifier(**self.lgbm_params).fit(X_fit, y) if len(classes) > 1 else None
+        fills = []
+        for X in X_miss:
+            if len(X) == 0:
+                fills.append(np.empty(0, dtype=np.int64))
+                continue
+            probs = np.zeros((len(X), K))
+            if clf is None:
+                probs[:, classes[0]] = 1.0
+            else:
+                probs[:, clf.classes_] = clf.predict_proba(X)
+            probs = (1 - self.alpha) * probs + self.alpha / K
+            u = rng.random(len(X))[:, None]
+            fills.append(np.minimum((probs.cumsum(axis=1) < u).sum(axis=1), K - 1))
+        return fills

augmask/baselines/__init__.py

@@ -0,0 +1 @@
+"""Missing-aware baselines: ForestDiff (official package) and MIWAE."""

augmask/baselines/forestdiff.py

@@ -0,0 +1,73 @@
+"""ForestDiff baseline: the official ForestDiffusion package with the paper's settings.
+
+Appendix C.6: n_t = 20 and duplicate_K = 10 (the defaults 50/100 did not finish
+within 3600 s); on News the training data is streamed in chunks of 8,000 rows.
+Everything else uses the package defaults (flow matching, XGBoost hyperparameters).
+The model is trained in data space on the incomplete train table (NaN = missing).
+
+One fix: the package one-hot encodes categoricals with >= 3 levels via
+pd.get_dummies, which turns a missing entry into the reference level. We keep
+those entries missing (NaN in every dummy), so they are skipped as targets and
+treated as missing inputs, as for numerical and binary columns.
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pandas as pd
+
+from ..data.schema import Schema
+
+CHUNK_SIZE = {"news": 8000}
+
+
+def _model_class():
+    from ForestDiffusion import ForestDiffusionModel
+
+    class _Model(ForestDiffusionModel):
+        def dummify(self, X):
+            out, before, after = super().dummify(X)
+            for i in self.cat_indexes:
+                dummies = [k for k, name in enumerate(after) if name.startswith(f"{i}_")]
+                out[np.ix_(np.isnan(X[:, i]), dummies)] = np.nan
+            return out, before, after
+
+    return _Model
+
+
+class ForestDiff:
+    def __init__(self, n_t: int = 20, duplicate_K: int = 10, chunk_size: int | None = None,
+                 n_jobs: int = -1, seed: int = 0, **forest_kwargs):
+        self.n_t, self.duplicate_K, self.chunk_size = n_t, duplicate_K, chunk_size
+        # one XGBoost thread per model: the package already runs n_jobs models in parallel
+        self.n_jobs, self.seed, self.forest_kwargs = n_jobs, seed, {"nthread": 1, **forest_kwargs}
+
+    def fit(self, table: pd.DataFrame, schema: Schema) -> "ForestDiff":
+        self.schema = schema
+        num, cat = list(schema.num_block), list(schema.cat_block)
+        self.vocab_ = [sorted(table[c].dropna().unique().tolist()) for c in cat]
+        codes = [table[c].map({v: i for i, v in enumerate(voc)}).to_numpy(np.float64)
+                 for c, voc in zip(cat, self.vocab_)]
+        X = np.column_stack([table[num].to_numpy(np.float64), *codes]) if codes else table[num].to_numpy(np.float64)
+        d_num = len(num)
+        K = [len(v) for v in self.vocab_]
+        int_idx = [j for j, c in enumerate(num) if (table[c].dropna() % 1 == 0).all()]  # rounded when sampling
+        bin_idx = [d_num + k for k, n in enumerate(K) if n <= 2]
+        cat_idx = [d_num + k for k, n in enumerate(K) if n > 2]
+        n_batch = math.ceil(len(X) / self.chunk_size) if self.chunk_size else 1  # the package default
+        self.model_ = _model_class()(
+            X, n_t=self.n_t, duplicate_K=self.duplicate_K, bin_indexes=bin_idx, cat_indexes=cat_idx, int_indexes=int_idx,
+            n_jobs=self.n_jobs, n_batch=n_batch, seed=self.seed, **self.forest_kwargs,
+        )
+        return self
+
+    def sample(self, n: int) -> pd.DataFrame:
+        np.random.seed(self.seed)
+        Xs = self.model_.generate(batch_size=n)
+        d_num = len(self.schema.num_block)
+        out = pd.DataFrame(Xs[:, :d_num], columns=list(self.schema.num_block))
+        for k, c in enumerate(self.schema.cat_block):
+            codes = np.clip(np.rint(Xs[:, d_num + k]).astype(np.int64), 0, len(self.vocab_[k]) - 1)
+            out[c] = np.asarray(self.vocab_[k], dtype=object)[codes]
+        return out[list(self.schema.columns)]

augmask/baselines/miwae.py

@@ -0,0 +1,110 @@
+"""MIWAE baseline (Mattei & Frellsen, 2019) used as a generator.
+
+Architecture and training follow the hyperimpute MIWAE imputer: two-layer ReLU
+encoder/decoder, Gaussian q(z|x), Student-t p(x|z) for numerical features, the
+importance-weighted bound over the observed entries, Adam (lr 1e-3), batch 256,
+zero-filled inputs. Categorical features get a softmax head and one-hot inputs.
+Settings from Appendix C.6: 800 epochs, latent size 64, hidden width 256, K = 40.
+Samples are generated by drawing z from the prior and x from p(x|z).
+
+Inputs and outputs live in the scenario's model space (see augmask.data.scenario).
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.distributions import Normal, StudentT
+
+
+
+def _mlp(d_in: int, h: int, d_out: int) -> nn.Sequential:
+    net = nn.Sequential(nn.Linear(d_in, h), nn.ReLU(), nn.Linear(h, h), nn.ReLU(), nn.Linear(h, d_out))
+    for m in net:
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+    return net
+
+
+class _MIWAENet(nn.Module):
+    def __init__(self, d_num: int, n_categories: list[int], latent: int, hidden: int):
+        super().__init__()
+        self.d_num, self.n_categories, self.latent = d_num, list(n_categories), latent
+        d_in = d_num + sum(n_categories)
+        self.encoder = _mlp(d_in, hidden, 2 * latent)
+        self.decoder = _mlp(latent, hidden, 3 * d_num + sum(n_categories))
+
+    def _inputs(self, x_num, x_cat, m_num, m_cat):
+        parts = [x_num * m_num]
+        for k, n in enumerate(self.n_categories):
+            parts.append(F.one_hot(x_cat[:, k].clamp(min=0), n).float() * m_cat[:, k:k + 1])
+        return torch.cat(parts, dim=1)
+
+    def _decode(self, z):
+        out = self.decoder(z)
+        d = self.d_num
+        loc, scale, df = out[..., :d], F.softplus(out[..., d:2 * d]) + 1e-3, F.softplus(out[..., 2 * d:3 * d]) + 3
+        logits = torch.split(out[..., 3 * d:], self.n_categories, dim=-1) if self.n_categories else ()
+        return loc, scale, df, logits
+
+    def neg_bound(self, x_num, x_cat, m_num, m_cat, K: int):
+        enc = self.encoder(self._inputs(x_num, x_cat, m_num, m_cat))
+        q = Normal(enc[:, :self.latent], F.softplus(enc[:, self.latent:]) + 1e-4)
+        z = q.rsample((K,))  # [K, B, latent]
+        loc, scale, df, logits = self._decode(z)
+        log_px = (StudentT(df, loc, scale).log_prob(x_num) * m_num).sum(-1)
+        for k, lg in enumerate(logits):
+            ll = torch.log_softmax(lg, -1).gather(-1, x_cat[:, k].clamp(min=0).expand(K, -1).unsqueeze(-1)).squeeze(-1)
+            log_px = log_px + ll * m_cat[:, k]
+        log_pz = Normal(0.0, 1.0).log_prob(z).sum(-1)
+        log_qz = q.log_prob(z).sum(-1)
+        return -(torch.logsumexp(log_px + log_pz - log_qz, dim=0) - math.log(K)).mean()
+
+    @torch.no_grad()
+    def sample(self, n: int, device):
+        z = torch.randn(n, self.latent, device=device)
+        loc, scale, df, logits = self._decode(z)
+        x_num = StudentT(df, loc, scale).sample()
+        x_cat = [torch.distributions.Categorical(logits=lg).sample() for lg in logits]
+        x_cat = torch.stack(x_cat, 1) if x_cat else torch.zeros(n, 0, dtype=torch.long, device=device)
+        return x_num, x_cat
+
+
+class MIWAE:
+    def __init__(self, n_epochs: int = 800, latent_size: int = 64, n_hidden: int = 256, K: int = 40,
+                 batch_size: int = 256, lr: float = 1e-3, seed: int = 0, device: str | None = None):
+        self.n_epochs, self.latent_size, self.n_hidden, self.K = n_epochs, latent_size, n_hidden, K
+        self.batch_size, self.lr, self.seed = batch_size, lr, seed
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+
+    def fit(self, X_num, X_cat, M_num, M_cat, n_categories) -> "MIWAE":
+        torch.manual_seed(self.seed)
+        rng = np.random.default_rng(self.seed)
+        dev = self.device
+        t = lambda a, dt: torch.as_tensor(np.asarray(a), dtype=dt, device=dev)
+        x_num, x_cat = t(np.where(M_num, X_num, 0.0), torch.float32), t(np.where(M_cat, X_cat, 0), torch.long)
+        m_num, m_cat = t(M_num, torch.float32), t(M_cat, torch.float32)
+        self.net_ = _MIWAENet(x_num.shape[1], list(n_categories), self.latent_size, self.n_hidden).to(dev)
+        opt = torch.optim.Adam(self.net_.parameters(), lr=self.lr)
+        n = len(x_num)
+        self.history_ = []
+        for epoch in range(self.n_epochs):
+            perm = torch.as_tensor(rng.permutation(n), device=dev)
+            total = torch.zeros((), device=dev)
+            for idx in torch.split(perm, min(self.batch_size, n)):
+                loss = self.net_.neg_bound(x_num[idx], x_cat[idx], m_num[idx], m_cat[idx], self.K)
+                opt.zero_grad()
+                loss.backward()
+                opt.step()
+                total += loss.detach() * len(idx)
+            self.history_.append({"epoch": epoch + 1, "neg_bound": (total / n).item()})
+        return self
+
+    def sample(self, n: int) -> tuple[np.ndarray, np.ndarray]:
+        torch.manual_seed(self.seed)
+        self.net_.eval()
+        x_num, x_cat = self.net_.sample(n, self.device)
+        return x_num.cpu().numpy(), x_cat.cpu().numpy()

augmask/data/__init__.py

@@ -0,0 +1,7 @@
+from .missingness import make_masks
+from .scenario import Scenario, build_scenario, load_artifact, save_artifact
+from .schema import Schema
+from .split import split_indices
+
+__all__ = ["Schema", "Scenario", "build_scenario", "load_artifact", "save_artifact",
+           "make_masks", "split_indices"]

augmask/data/datasets.py

@@ -0,0 +1,190 @@
+"""Step 1: download and clean the six UCI benchmark tables.
+
+Each loader returns the complete table (no NaN) and its Schema. Numerical
+columns are float64; categorical columns and a classification target are str.
+Native missing categories (Adult's "?") are kept as their own category; rows
+with a missing numerical value or target are dropped (Beijing's pm2.5).
+"""
+from __future__ import annotations
+
+import hashlib
+import os
+import shutil
+import urllib.request
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from .schema import Schema
+
+DATASETS = ("adult", "default", "shoppers", "magic", "beijing", "news")
+
+# archive URL and sha256 of the data file(s) used inside it
+SOURCES = {
+    "adult": ("https://archive.ics.uci.edu/static/public/2/adult.zip", {
+        "adult.data": "5b00264637dbfec36bdeaab5676b0b309ff9eb788d63554ca0a249491c86603d",
+        "adult.test": "a2a9044bc167a35b2361efbabec64e89d69ce82d9790d2980119aac5fd7e9c05"}),
+    "default": ("https://archive.ics.uci.edu/static/public/350/default+of+credit+card+clients.zip", {
+        "default of credit card clients.xls": "30c6be3abd8dcfd3e6096c828bad8c2f011238620f5369220bd60cfc82700933"}),
+    "shoppers": ("https://archive.ics.uci.edu/static/public/468/online+shoppers+purchasing+intention+dataset.zip", {
+        "online_shoppers_intention.csv": "b3055ee355f59134d851d32641183cb4a8b45def7124d2f50442a042f358e0d9"}),
+    "magic": ("https://archive.ics.uci.edu/static/public/159/magic+gamma+telescope.zip", {
+        "magic04.data": "e9314b7ebd4b4b59a3b3d65f7316663963777b16a46786877651dbbaa640b36a"}),
+    "beijing": ("https://archive.ics.uci.edu/static/public/381/beijing+pm2+5+data.zip", {
+        "PRSA_data_2010.1.1-2014.12.31.csv": "4127f868775e31b3956522adc0ec75af8937dde6a3896e8beed3a376c6d27f1c"}),
+    "news": ("https://archive.ics.uci.edu/static/public/332/online+news+popularity.zip", {
+        "OnlineNewsPopularity/OnlineNewsPopularity.csv":
+            "b66d9088632308cc27fa35af847650d174a5a50503987c4e511de94a99d1c218"}),
+}
+
+
+def data_root() -> Path:
+    return Path(os.environ.get("AUGMASK_DATA", "datasets"))
+
+
+def load_dataset(name: str, root: str | Path | None = None, verify: bool = True) -> tuple[pd.DataFrame, Schema]:
+    """Complete table and schema of ``name``; downloads and caches on first use."""
+    if name not in DATASETS:
+        raise ValueError(f"unknown dataset {name!r}; choose from {DATASETS}")
+    root = Path(root) if root is not None else data_root()
+    out = root / name
+    if (out / "complete.csv").exists() and (out / "schema.json").exists():
+        schema = Schema.from_json(out / "schema.json")
+        df = pd.read_csv(out / "complete.csv", dtype={c: str for c in schema.cat_block}, keep_default_na=False)
+        return _finalize(df, schema), schema
+    raw = _download(name, root / "raw", verify)
+    df, schema = LOADERS[name](raw)
+    df = _finalize(df, schema)
+    out.mkdir(parents=True, exist_ok=True)
+    df.to_csv(out / "complete.csv", index=False)
+    schema.to_json(out / "schema.json")
+    return df, schema
+
+
+def _finalize(df: pd.DataFrame, schema: Schema) -> pd.DataFrame:
+    df = df[list(schema.columns)].reset_index(drop=True)
+    for c in schema.num_block:
+        df[c] = df[c].astype(np.float64)
+    for c in schema.cat_block:
+        df[c] = df[c].astype(str)
+    if df.isna().any().any():
+        raise ValueError(f"{schema.name}: complete table contains NaN")
+    return df
+
+
+def _sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def _download(name: str, raw_root: Path, verify: bool) -> Path:
+    url, files = SOURCES[name]
+    dest = raw_root / name
+    if not all((dest / f).exists() for f in files):
+        dest.mkdir(parents=True, exist_ok=True)
+        archive = dest / "archive.zip"
+        with urllib.request.urlopen(url) as r, open(archive, "wb") as f:
+            shutil.copyfileobj(r, f)
+        _safe_extract(archive, dest)
+        for nested in list(dest.rglob("*.zip")):
+            if nested != archive:
+                _safe_extract(nested, nested.parent)
+    if verify:
+        for f, digest in files.items():
+            if _sha256(dest / f) != digest:
+                raise ValueError(f"{name}: {f} does not match the expected sha256; the source may have changed")
+    return dest
+
+
+def _safe_extract(archive: Path, dest: Path) -> None:
+    root = dest.resolve()
+    with zipfile.ZipFile(archive) as z:
+        for member in z.namelist():
+            if not (root / member).resolve().is_relative_to(root):
+                raise ValueError(f"unsafe path in archive: {member}")
+        z.extractall(root)
+
+
+def _cat_str(s: pd.Series) -> pd.Series:
+    """Categories as strings; integer-valued codes without a trailing '.0'."""
+    if pd.api.types.is_bool_dtype(s):
+        return s.map({True: "True", False: "False"})
+    if pd.api.types.is_numeric_dtype(s):
+        return s.astype(np.int64).astype(str)
+    return s.astype(str).str.strip()
+
+
+# ----------------------------------------------------------------------------- loaders
+def _adult(raw: Path):
+    cols = ["age", "workclass", "fnlwgt", "education", "education.num", "marital.status", "occupation",
+            "relationship", "race", "sex", "capital.gain", "capital.loss", "hours.per.week",
+            "native.country", "income"]
+    read = dict(header=None, names=cols, skipinitialspace=True, keep_default_na=False)
+    df = pd.concat([pd.read_csv(raw / "adult.data", **read),
+                    pd.read_csv(raw / "adult.test", skiprows=1, **read)], ignore_index=True)
+    df["income"] = df["income"].str.rstrip(".")
+    num = ("age", "fnlwgt", "education.num", "capital.gain", "capital.loss", "hours.per.week")
+    cat = tuple(c for c in cols[:-1] if c not in num)
+    for c in cat + ("income",):
+        df[c] = _cat_str(df[c])
+    return df, Schema("adult", "binclass", tuple(cols), num, cat, "income")
+
+
+def _default(raw: Path):
+    df = pd.read_excel(raw / "default of credit card clients.xls", sheet_name="Data", header=1).drop(columns="ID")
+    target = "default payment next month"
+    cat = ("SEX", "EDUCATION", "MARRIAGE", "PAY_0", "PAY_2", "PAY_3", "PAY_4", "PAY_5", "PAY_6")
+    num = tuple(c for c in df.columns if c not in cat + (target,))
+    for c in cat + (target,):
+        df[c] = _cat_str(df[c])
+    return df, Schema("default", "binclass", tuple(df.columns), num, cat, target)
+
+
+def _shoppers(raw: Path):
+    df = pd.read_csv(raw / "online_shoppers_intention.csv")
+    cat = ("Month", "OperatingSystems", "Browser", "Region", "TrafficType", "VisitorType", "Weekend")
+    num = tuple(df.columns[:10])
+    for c in cat + ("Revenue",):
+        df[c] = _cat_str(df[c])
+    return df, Schema("shoppers", "binclass", tuple(df.columns), num, cat, "Revenue")
+
+
+def _magic(raw: Path):
+    cols = ["Length", "Width", "Size", "Conc", "Conc1", "Asym", "M3Long", "M3Trans", "Alpha", "Dist", "class"]
+    df = pd.read_csv(raw / "magic04.data", header=None, names=cols)
+    df["class"] = _cat_str(df["class"])
+    return df, Schema("magic", "binclass", tuple(cols), tuple(cols[:-1]), (), "class")
+
+
+def _beijing(raw: Path):
+    df = pd.read_csv(raw / "PRSA_data_2010.1.1-2014.12.31.csv").drop(columns="No")
+    df = df.dropna(subset=["pm2.5"]).reset_index(drop=True)
+    cat = ("year", "month", "day", "hour", "cbwd")
+    num = ("DEWP", "TEMP", "PRES", "Iws", "Is", "Ir")
+    for c in cat:
+        df[c] = _cat_str(df[c])
+    return df, Schema("beijing", "regression", tuple(df.columns), num, cat, "pm2.5")
+
+
+def _news(raw: Path):
+    df = pd.read_csv(raw / "OnlineNewsPopularity" / "OnlineNewsPopularity.csv", skipinitialspace=True)
+    df.columns = [c.strip() for c in df.columns]
+    channels = [c for c in df.columns if c.startswith("data_channel_is_")]
+    days = [c for c in df.columns if c.startswith("weekday_is_")]
+    onehot = df[channels].to_numpy()
+    df["data_channel"] = np.where(onehot.sum(1) == 0, "none",
+                                  np.array([c.removeprefix("data_channel_is_") for c in channels])[onehot.argmax(1)])
+    df["weekday"] = np.array([d.removeprefix("weekday_is_") for d in days])[df[days].to_numpy().argmax(1)]
+    df = df.drop(columns=["url", "is_weekend", *channels, *days])
+    num = tuple(c for c in df.columns if c not in ("shares", "data_channel", "weekday"))
+    cols = num + ("shares", "data_channel", "weekday")
+    return df[list(cols)], Schema("news", "regression", cols, num, ("data_channel", "weekday"), "shares")
+
+
+LOADERS = {"adult": _adult, "default": _default, "shoppers": _shoppers,
+           "magic": _magic, "beijing": _beijing, "news": _news}