Skills › Scientific Agent Skills › pytdc › load_and_split_data.py
load_and_split_data.py
399 lines · skills/pytdc/scripts/load_and_split_data.py at 5f2bc03
1#!/usr/bin/env python3
2"""Plan or explicitly execute one bounded PyTDC dataset split."""
4from __future__ import annotations
6import argparse
7import importlib
8import itertools
9import sys
10from pathlib import Path
11from typing import Any
13from _common import (
14 CliError,
15 bounded_int,
16 canonical_name,
17 emit_json,
18 load_pytdc_metadata,
19 safe_directory,
20 truncate_value,
21 validate_fractions,
22)
25# Public imports verified against the PyTDC 1.1.15 source distribution.
26TASKS: dict[str, tuple[str, str, str]] = {
27 "ADME": ("tdc.single_pred", "ADME", "ADME"),
28 "CRISPROutcome": ("tdc.single_pred", "CRISPROutcome", "CRISPROutcome"),
29 "Develop": ("tdc.single_pred", "Develop", "Develop"),
30 "Epitope": ("tdc.single_pred", "Epitope", "Epitope"),
31 "HTS": ("tdc.single_pred", "HTS", "HTS"),
32 "Paratope": ("tdc.single_pred", "Paratope", "Paratope"),
33 "QM": ("tdc.single_pred", "QM", "QM"),
34 "Tox": ("tdc.single_pred", "Tox", "Tox"),
35 "Yields": ("tdc.single_pred", "Yields", "Yields"),
36 "AntibodyAff": ("tdc.multi_pred", "AntibodyAff", "AntibodyAff"),
37 "Catalyst": ("tdc.multi_pred", "Catalyst", "Catalyst"),
38 "DDI": ("tdc.multi_pred", "DDI", "DDI"),
39 "DrugRes": ("tdc.multi_pred", "DrugRes", "DrugRes"),
40 "DrugSyn": ("tdc.multi_pred", "DrugSyn", "DrugSyn"),
41 "DTI": ("tdc.multi_pred", "DTI", "DTI"),
42 "GDA": ("tdc.multi_pred", "GDA", "GDA"),
43 "MTI": ("tdc.multi_pred", "MTI", "MTI"),
44 "PeptideMHC": ("tdc.multi_pred", "PeptideMHC", "PeptideMHC"),
45 "PPI": ("tdc.multi_pred", "PPI", "PPI"),
46 "ProteinPeptide": ("tdc.multi_pred", "ProteinPeptide", "ProteinPeptide"),
47 "TCREpitopeBinding": (
48 "tdc.multi_pred",
49 "TCREpitopeBinding",
50 "TCREpitopeBinding",
51 ),
52 "TrialOutcome": ("tdc.multi_pred", "TrialOutcome", "TrialOutcome"),
53 "MolGen": ("tdc.generation", "MolGen", "MolGen"),
54 "Reaction": ("tdc.generation", "Reaction", "Reaction"),
55 "RetroSyn": ("tdc.generation", "RetroSyn", "RetroSyn"),
56}
58SPLIT_METHODS = ("random", "scaffold", "cold_split", "combination", "time")
61def normalize_columns(values: list[str] | None) -> list[str]:
62 columns: list[str] = []
63 for value in values or []:
64 columns.extend(part.strip() for part in value.split(",") if part.strip())
65 columns = list(dict.fromkeys(columns))
66 if len(columns) > 8:
67 raise CliError("at most eight cold-split columns may be specified")
68 return columns
71def validate_label(task: str, dataset: str, label: str | None) -> str | None:
72 """Resolve multi-label targets from the download-free upstream registry."""
73 from tdc.utils import dataset2target_lists
75 choices = dataset2target_lists.get(dataset, [])
76 if choices:
77 if TASKS[task][0] == "tdc.generation" or task in {"DrugSyn", "TCREpitopeBinding"}:
78 raise CliError(f"{task}/{dataset} has no supported label-selecting constructor")
79 if label is None:
80 raise CliError(f"{dataset} requires --label-name; available: {', '.join(choices[:20])}")
81 return canonical_name(label, choices, "label")
82 if label is not None:
83 raise CliError(f"{dataset} has no target list in the reviewed package registry")
84 return None
87def validate_request(
88 *,
89 task_query: str,
90 dataset_query: str,
91 method: str,
92 columns: list[str],
93 time_column: str | None,
94 metadata: Any,
95) -> tuple[str, str]:
96 """Resolve exact registry names and reject unsupported split combinations."""
98 task = canonical_name(task_query, TASKS, "task")
99 module_name, _, registry_key = TASKS[task]
100 registry_key = canonical_name(registry_key, metadata.dataset_names, "task registry")
101 dataset = canonical_name(
102 dataset_query, metadata.dataset_names[registry_key], f"{task} dataset"
103 )
105 if method == "scaffold" and task not in {"ADME", "Tox", "HTS"}:
106 raise CliError(
107 "the official generic scaffold-split documentation limits this method "
108 "to the molecule-based ADME, Tox, and HTS single-instance tasks"
109 )
110 if method == "cold_split":
111 if module_name != "tdc.multi_pred":
112 raise CliError("cold_split is only exposed by multi-instance loaders")
113 if not columns:
114 raise CliError("cold_split requires one or more --column values")
115 elif columns:
116 raise CliError("--column is only valid with --method cold_split")
118 if method == "combination" and task != "DrugSyn":
119 raise CliError(
120 "the built-in combination split is documented for DrugSyn combination data"
121 )
123 if method == "time":
124 if not (
125 task == "DTI"
126 and dataset.casefold() == "bindingdb_patent"
127 and time_column
128 ):
129 raise CliError(
130 "the verified built-in temporal case is DTI/BindingDB_Patent and "
131 "requires --time-column (normally Year)"
132 )
133 elif time_column:
134 raise CliError("--time-column is only valid with --method time")
136 if module_name == "tdc.generation" and method != "random":
137 raise CliError("the verified MolGen/Reaction/RetroSyn loaders expose random split")
138 return task, dataset
141def build_plan(
142 *,
143 task: str,
144 dataset: str,
145 method: str,
146 seed: int,
147 fractions: tuple[float, float, float],
148 columns: list[str],
149 time_column: str | None,
150 data_dir: Path,
151 package_version: str,
152 label_name: str | None = None,
153) -> dict[str, Any]:
154 return {
155 "action": "plan",
156 "acknowledgement_required": "--execute",
157 "data_directory": str(data_dir),
158 "dataset": dataset,
159 "label_name": label_name,
160 "download_performed": False,
161 "network_and_storage": (
162 "Constructing the loader may contact TDC/Harvard Dataverse and write "
163 "the complete dataset under data_directory."
164 ),
165 "package": "PyTDC",
166 "package_version": package_version,
167 "split": {
168 "cold_columns": columns,
169 "fractions": list(fractions),
170 "method": method,
171 "seed": seed,
172 "time_column": time_column,
173 },
174 "task": task,
175 }
178def _summarize_frame(frame: Any, preview: int) -> dict[str, Any]:
179 summary: dict[str, Any] = {"rows": len(frame)}
180 if hasattr(frame, "columns"):
181 summary["columns"] = [str(column) for column in list(frame.columns)[:100]]
182 if preview and hasattr(frame, "head") and hasattr(frame, "to_dict"):
183 records = frame.head(preview).to_dict(orient="records")
184 summary["preview"] = truncate_value(records)
185 return summary
188def _audit_cold_columns(
189 split: dict[str, Any], columns: list[str]
190) -> dict[str, Any]:
191 """Report exact-value overlap; this is not a general leakage proof."""
193 audit: dict[str, Any] = {}
194 partitions = ("train", "valid", "test")
195 for column in columns:
196 missing = [
197 name
198 for name in partitions
199 if not hasattr(split[name], "columns") or column not in split[name].columns
200 ]
201 if missing:
202 audit[column] = {"missing_from": missing}
203 continue
204 values = {
205 name: set(str(value) for value in split[name][column].dropna().unique())
206 for name in partitions
207 }
208 audit[column] = {
209 "distinct_values": {name: len(items) for name, items in values.items()},
210 "pairwise_overlap_counts": {
211 f"{left}:{right}": len(values[left] & values[right])
212 for left, right in itertools.combinations(partitions, 2)
213 },
214 }
215 return audit
218def execute_split(
219 *,
220 task: str,
221 dataset: str,
222 method: str,
223 seed: int,
224 fractions: tuple[float, float, float],
225 columns: list[str],
226 time_column: str | None,
227 data_dir: Path,
228 preview: int,
229 package_version: str,
230 label_name: str | None = None,
231) -> dict[str, Any]:
232 """Instantiate a loader only after the caller acknowledges the download."""
234 module_name, class_name, _ = TASKS[task]
235 task_class = getattr(importlib.import_module(module_name), class_name)
236 constructor_kwargs = {"label_name": label_name} if label_name is not None else {}
237 loader = task_class(name=dataset, path=str(data_dir), **constructor_kwargs)
239 split_kwargs: dict[str, Any] = {
240 "method": method,
241 "seed": seed,
242 "frac": list(fractions),
243 }
244 if method == "cold_split":
245 split_kwargs["column_name"] = columns
246 if method == "time":
247 split_kwargs["time_column"] = time_column
248 split = loader.get_split(**split_kwargs)
250 expected = {"train", "valid", "test"}
251 if not isinstance(split, dict) or not expected.issubset(split):
252 raise CliError("PyTDC returned an unexpected split structure")
254 result: dict[str, Any] = {
255 "action": "executed",
256 "data_directory": str(data_dir),
257 "dataset": dataset,
258 "label_name": label_name,
259 "download_acknowledged": True,
260 "package": "PyTDC",
261 "package_version": package_version,
262 "partitions": {
263 name: _summarize_frame(split[name], preview)
264 for name in ("train", "valid", "test")
265 },
266 "split": {
267 "cold_columns": columns,
268 "fractions": list(fractions),
269 "method": method,
270 "seed": seed,
271 "time_column": time_column,
272 },
273 "task": task,
274 }
275 if method == "cold_split":
276 result["cold_column_audit"] = _audit_cold_columns(split, columns)
277 result["cold_column_audit_note"] = (
278 "Counts cover exact values in requested columns only; zero overlap is "
279 "not a blanket claim that all biological or chemical leakage is absent."
280 )
281 return result
284def build_parser() -> argparse.ArgumentParser:
285 parser = argparse.ArgumentParser(
286 description=(
287 "Validate and plan a PyTDC dataset split. The default is download-free; "
288 "pass --execute to acknowledge that loader construction may download "
289 "and cache the full dataset."
290 )
291 )
292 parser.add_argument("--task", required=True, help="exact public PyTDC task class")
293 parser.add_argument("--dataset", required=True, help="exact package-registry name")
294 parser.add_argument("--label-name", help="required target for a multi-label dataset")
295 parser.add_argument(
296 "--method",
297 choices=SPLIT_METHODS,
298 default="random",
299 help="split method (default: random)",
300 )
301 parser.add_argument(
302 "--column",
303 action="append",
304 help="cold-split column; repeat or pass comma-separated values",
305 )
306 parser.add_argument("--time-column", help="time column for --method time")
307 parser.add_argument(
308 "--frac",
309 nargs=3,
310 type=float,
311 metavar=("TRAIN", "VALID", "TEST"),
312 default=(0.7, 0.1, 0.2),
313 help="split fractions summing to 1 (default: 0.7 0.1 0.2)",
314 )
315 parser.add_argument(
316 "--seed",
317 type=bounded_int(0, 4_294_967_295),
318 default=42,
319 help="non-negative split seed (default: 42)",
320 )
321 parser.add_argument(
322 "--data-dir",
323 default=".pytdc-data",
324 help="relative cache directory (default: .pytdc-data)",
325 )
326 parser.add_argument(
327 "--preview",
328 type=bounded_int(0, 10),
329 default=0,
330 help="include at most this many bounded rows per partition (default: 0)",
331 )
332 parser.add_argument(
333 "--execute",
334 action="store_true",
335 help="acknowledge network/storage use and construct the loader",
336 )
337 parser.add_argument("--output", help="write JSON to a relative workspace path")
338 parser.add_argument(
339 "--force", action="store_true", help="replace an existing --output file"
340 )
341 return parser
344def main(argv: list[str] | None = None) -> int:
345 args = build_parser().parse_args(argv)
346 try:
347 fractions = validate_fractions(args.frac)
348 columns = normalize_columns(args.column)
349 metadata, package_version = load_pytdc_metadata()
350 task, dataset = validate_request(
351 task_query=args.task,
352 dataset_query=args.dataset,
353 method=args.method,
354 columns=columns,
355 time_column=args.time_column,
356 metadata=metadata,
357 )
358 label_name = validate_label(task, dataset, args.label_name)
359 data_dir = safe_directory(
360 args.data_dir,
361 label="data directory",
362 create=args.execute,
363 )
364 if args.execute:
365 result = execute_split(
366 task=task,
367 dataset=dataset,
368 method=args.method,
369 seed=args.seed,
370 fractions=fractions,
371 columns=columns,
372 time_column=args.time_column,
373 data_dir=data_dir,
374 preview=args.preview,
375 package_version=package_version,
376 label_name=label_name,
377 )
378 else:
379 result = build_plan(
380 task=task,
381 dataset=dataset,
382 method=args.method,
383 seed=args.seed,
384 fractions=fractions,
385 columns=columns,
386 time_column=args.time_column,
387 data_dir=data_dir,
388 package_version=package_version,
389 label_name=label_name,
390 )
391 emit_json(result, args.output, force=args.force)
392 except (CliError, ImportError, OSError, TypeError, ValueError) as exc:
393 print(f"error: {exc}", file=sys.stderr)
394 return 2
395 return 0
398if __name__ == "__main__":
399 raise SystemExit(main())