"""Dataset fetching for the reproductions. Vendors nothing, pins everything. Two datasets, both from their canonical public hosts, both checksummed. The checksums are the point: a reproduction whose inputs are unverified is a reproduction of whatever the mirror served that day. Files land in ``~/.cache/torch_dimensions`` (override with ``TD_DATA_DIR``) and are downloaded once. No torchvision. The library depends on ``torch`` and nothing else, and a reproduction that quietly needs a second framework is not reproducing with the library — it is reproducing with the library plus whatever else was lying around. """ from __future__ import annotations import gzip import hashlib import os import tarfile from pathlib import Path from urllib.request import urlopen import torch MNIST_URL = "https://ossci-datasets.s3.amazonaws.com/mnist" MNIST_FILES = { "train-images-idx3-ubyte.gz": ( "440fcabf73cc546fa21475e81ea370265605f56be210a4024d2ca8f203523609" ), "train-labels-idx1-ubyte.gz": ( "3552534a0a558bbed6aed32b30c495cca23d567ec52cac8be1a0730e8010255c" ), "t10k-images-idx3-ubyte.gz": ( "8d422c7b0a1c1c79245a5bcf07fe86e33eeafee792b84584aec276f5a2dbc4e6" ), "t10k-labels-idx1-ubyte.gz": ( "f7ae60f92e00ec6debd23a6088c31dbd2371eca3ffa0defaefb259924204aec6" ), } CIFAR_URL = "https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz" CIFAR_SHA = "c4a38c50a1bc5f3a1c5537f2155ab9d68f9f25eb1ed8d9ddda3db29a59bca1dd" # Beijing Multi-Site Air-Quality (UCI 501): 12 monitoring stations x 6 # pollutants x 35,064 hourly steps. Chosen for the sparse-lattice work because # it is genuinely 2-D — station and pollutant are different kinds of axis, not # a reshaped sequence — and because its gaps are real measurement gaps. BEIJING_URL = ( "https://archive.ics.uci.edu/static/public/501/beijing+multi+site+air+quality+data.zip" ) BEIJING_SHA = "b04da438b2f331ac0ffd45aebdfec0d20d2367feb5f6948c4b1f7ce1191e33c4" POLLUTANTS = ("PM2.5", "PM10", "SO2", "NO2", "CO", "O3") def data_dir() -> Path: root = Path(os.environ.get("TD_DATA_DIR", Path.home() / ".cache" / "torch_dimensions")) root.mkdir(parents=True, exist_ok=True) return root def fetch(url: str, dest: Path, sha256: str | None) -> Path: """Download once, verify always. The checksum is re-checked on every call, not only after a download: a truncated or half-written cache file is exactly the failure that presents as "the model suddenly stopped learning". """ dest.parent.mkdir(parents=True, exist_ok=True) if not dest.exists(): print(f"downloading {url} -> {dest}") with urlopen(url) as r, open(dest, "wb") as fh: # noqa: S310 — pinned https URLs while chunk := r.read(1 << 20): fh.write(chunk) if sha256: got = hashlib.sha256(dest.read_bytes()).hexdigest() if got != sha256: raise RuntimeError( f"checksum mismatch for {dest.name}\n expected {sha256}\n got {got}\n" "delete the file and retry; if it repeats, the upstream host changed and this " "reproduction should not be trusted until the pin is reviewed" ) return dest def _idx(path: Path) -> torch.Tensor: """Read an IDX file (the MNIST container format) into a uint8 tensor.""" with gzip.open(path, "rb") as fh: raw = fh.read() if raw[0:2] != b"\x00\x00" or raw[2] != 0x08: raise ValueError(f"{path.name}: not a uint8 IDX file") n_dims = raw[3] dims = [int.from_bytes(raw[4 + 4 * i : 8 + 4 * i], "big") for i in range(n_dims)] body = torch.frombuffer(bytearray(raw[4 + 4 * n_dims :]), dtype=torch.uint8) return body.reshape(*dims) def mnist() -> dict[str, torch.Tensor]: """MNIST as raw uint8 tensors: images ``(N, 28, 28)``, labels ``(N,)``.""" root = data_dir() / "mnist" for name, sha in MNIST_FILES.items(): fetch(f"{MNIST_URL}/{name}", root / name, sha) return { "train_x": _idx(root / "train-images-idx3-ubyte.gz"), "train_y": _idx(root / "train-labels-idx1-ubyte.gz"), "test_x": _idx(root / "t10k-images-idx3-ubyte.gz"), "test_y": _idx(root / "t10k-labels-idx1-ubyte.gz"), } def cifar10() -> dict[str, torch.Tensor]: """CIFAR-10 as raw uint8: images ``(N, 32, 32, 3)``, labels ``(N,)``.""" root = data_dir() / "cifar10" archive = fetch(CIFAR_URL, root / "cifar-10-binary.tar.gz", CIFAR_SHA) members = { "train": [f"cifar-10-batches-bin/data_batch_{i}.bin" for i in range(1, 6)], "test": ["cifar-10-batches-bin/test_batch.bin"], } out: dict[str, torch.Tensor] = {} with tarfile.open(archive) as tar: for split, names in members.items(): blobs = [] for name in names: fh = tar.extractfile(name) if fh is None: raise RuntimeError(f"{archive.name} is missing {name}") blobs.append(torch.frombuffer(bytearray(fh.read()), dtype=torch.uint8)) # Each record is 1 label byte followed by 3072 bytes of plane-major # RGB — 1024 red, then green, then blue. recs = torch.cat(blobs).reshape(-1, 3073) out[f"{split}_y"] = recs[:, 0].clone() out[f"{split}_x"] = recs[:, 1:].reshape(-1, 3, 32, 32).permute(0, 2, 3, 1).contiguous() return out def beijing() -> dict: """Air quality as a ``(T, station, pollutant)`` tensor plus its axis names. Returns the raw readings with gaps marked as NaN — imputation is the experiment's decision, not the loader's, and a loader that quietly fills holes is a loader that decides how missing data behaves for everyone downstream. """ import csv import io import zipfile root = data_dir() / "beijing" archive = fetch(BEIJING_URL, root / "beijing.zip", BEIJING_SHA) outer = zipfile.ZipFile(archive) inner_name = next(n for n in outer.namelist() if n.endswith(".zip")) inner = zipfile.ZipFile(io.BytesIO(outer.read(inner_name))) csv_names = sorted(n for n in inner.namelist() if n.endswith(".csv")) stations, columns = [], [] for name in csv_names: rows = list(csv.DictReader(io.StringIO(inner.read(name).decode()))) stations.append(rows[0]["station"]) series = torch.full((len(rows), len(POLLUTANTS)), float("nan")) for t, row in enumerate(rows): for j, pollutant in enumerate(POLLUTANTS): value = row[pollutant] if value not in ("NA", "", None): series[t, j] = float(value) columns.append(series) n = min(s.shape[0] for s in columns) series = torch.stack([s[:n] for s in columns], dim=1) # (T, station, pollutant) return { "series": series, "names": ("station", "pollutant"), "stations": tuple(stations), "pollutants": POLLUTANTS, } if __name__ == "__main__": # `python -m examples.repro.data` pre-fetches everything m = mnist() print("mnist", {k: tuple(v.shape) for k, v in m.items()}) c = cifar10() print("cifar10", {k: tuple(v.shape) for k, v in c.items()}) b = beijing() print("beijing", tuple(b["series"].shape), b["stations"][:3])