Skills › Scientific Agent Skills › pytdc › molecular_generation.py
molecular_generation.py
434 lines · skills/pytdc/scripts/molecular_generation.py at 5f2bc03
1#!/usr/bin/env python3
2"""Safely plan PyTDC molecule data access or bounded oracle scoring."""
4from __future__ import annotations
6import argparse
7import contextlib
8import os
9import sys
10from pathlib import Path
11from typing import Any, Iterator
13from _common import (
14 CliError,
15 bounded_int,
16 canonical_name,
17 emit_json,
18 load_pytdc_metadata,
19 safe_directory,
20 safe_input_file,
21 truncate_value,
22 validate_fractions,
23)
26LOCAL_SCALAR_ORACLES = {"qed"}
27# LogP and SA are listed as "trivial" upstream, but their 1.1.15 implementation
28# lazily downloads the fpscores artifact through calculateScore().
29CHECKPOINT_ORACLES = {
30 "logp",
31 "sa",
32 "drd2",
33 "gsk3b",
34 "jnk3",
35 "cyp3a4_veith",
36}
37MAX_SMILES_FILE_BYTES = 1 * 1024 * 1024
40@contextlib.contextmanager
41def _working_directory(path: Path) -> Iterator[None]:
42 original = Path.cwd()
43 os.chdir(path)
44 try:
45 yield
46 finally:
47 os.chdir(original)
50def classify_oracle(metadata: Any, query: str) -> tuple[str, str]:
51 name = canonical_name(query, metadata.oracle_names, "oracle")
52 if name in LOCAL_SCALAR_ORACLES:
53 return name, "local_scalar"
54 if name in CHECKPOINT_ORACLES:
55 return name, "checkpoint_download"
56 if name in set(metadata.download_oracle_names):
57 return name, "unsupported_checkpoint"
58 if name in set(metadata.distribution_oracles):
59 return name, "distribution_metric"
60 if name in set(metadata.download_receptor_oracle_name):
61 return name, "receptor_or_docking"
62 if name in set(metadata.synthetic_oracle_name):
63 return name, "remote_service"
64 return name, "specialized_or_composite"
67def load_smiles(
68 direct: list[str] | None,
69 input_path: str | None,
70 *,
71 max_molecules: int,
72) -> list[str]:
73 values = list(direct or [])
74 if input_path:
75 path = safe_input_file(
76 input_path,
77 max_bytes=MAX_SMILES_FILE_BYTES,
78 label="SMILES input",
79 )
80 with path.open("r", encoding="utf-8") as handle:
81 values.extend(line.strip() for line in handle if line.strip())
82 if not values:
83 raise CliError("provide at least one --smiles or --input line")
84 if len(values) > max_molecules:
85 raise CliError(
86 f"received {len(values)} molecules; --max-molecules is {max_molecules}"
87 )
88 if any(len(value) > 4096 for value in values):
89 raise CliError("each SMILES string is limited to 4096 characters")
90 return values
93def score_plan(
94 *,
95 oracle: str,
96 category: str,
97 smiles_count: int,
98 runtime_dir: Path,
99 package_version: str,
100 download_acknowledged: bool,
101) -> dict[str, Any]:
102 return {
103 "action": "plan",
104 "acknowledgement_required": (
105 "--execute"
106 if category == "local_scalar"
107 else "--execute --download (only supported checkpoint oracles)"
108 ),
109 "download_acknowledged": download_acknowledged,
110 "download_performed": False,
111 "oracle": oracle,
112 "oracle_category": category,
113 "package": "PyTDC",
114 "package_version": package_version,
115 "runtime_directory": str(runtime_dir),
116 "score_direction": "not assumed; consult the exact oracle documentation",
117 "smiles_count": smiles_count,
118 }
121def execute_scores(
122 *,
123 oracle_name: str,
124 category: str,
125 smiles: list[str],
126 runtime_dir: Path,
127 download_acknowledged: bool,
128 package_version: str,
129) -> dict[str, Any]:
130 if category == "checkpoint_download" and not download_acknowledged:
131 raise CliError(
132 f"{oracle_name} may download a model checkpoint; pass --download "
133 "together with --execute to acknowledge this"
134 )
135 if category not in {"local_scalar", "checkpoint_download"}:
136 raise CliError(
137 "this helper executes only local QED or the explicitly acknowledged "
138 "LogP/SA/DRD2/GSK3B/JNK3/CYP3A4_Veith checkpoint-backed oracles; "
139 "remote services, docking, distribution, and composite oracles are "
140 "intentionally not called"
141 )
143 from tdc import Oracle # Lazy optional import.
144 from rdkit import Chem
146 # Empty SMILES parses as a zero-atom molecule; it is not a scorable structure.
147 valid = []
148 for value in smiles:
149 molecule = Chem.MolFromSmiles(value)
150 valid.append(molecule is not None and molecule.GetNumAtoms() > 0)
151 valid_smiles = [value for value, ok in zip(smiles, valid) if ok]
153 with _working_directory(runtime_dir):
154 if valid_smiles:
155 oracle = Oracle(name=oracle_name, num_max_call=len(valid_smiles))
156 scores = oracle(valid_smiles)
157 else:
158 scores = []
159 if not isinstance(scores, list) or len(scores) != len(valid_smiles):
160 raise CliError("PyTDC returned an unexpected oracle result shape")
161 score_iter = iter(scores)
163 return {
164 "action": "executed",
165 "download_acknowledged": download_acknowledged,
166 "oracle": oracle_name,
167 "oracle_category": category,
168 "package": "PyTDC",
169 "package_version": package_version,
170 "results": [
171 {
172 "index": index,
173 "valid": ok,
174 "score": truncate_value(next(score_iter)) if ok else None,
175 "smiles": truncate_value(smiles_value),
176 }
177 for index, (smiles_value, ok) in enumerate(zip(smiles, valid))
178 ],
179 "runtime_directory": str(runtime_dir),
180 "score_direction": "not assumed; results preserve input order",
181 }
184def dataset_plan(
185 *,
186 dataset: str,
187 data_dir: Path,
188 seed: int,
189 fractions: tuple[float, float, float],
190 package_version: str,
191 download_acknowledged: bool,
192) -> dict[str, Any]:
193 return {
194 "action": "plan",
195 "acknowledgement_required": "--execute --download",
196 "data_directory": str(data_dir),
197 "dataset": dataset,
198 "download_acknowledged": download_acknowledged,
199 "download_performed": False,
200 "network_and_storage": (
201 "MolGen construction may download a complete, potentially very large "
202 "molecule corpus to data_directory."
203 ),
204 "package": "PyTDC",
205 "package_version": package_version,
206 "split": {
207 "fractions": list(fractions),
208 "method": "random",
209 "seed": seed,
210 },
211 }
214def execute_dataset(
215 *,
216 dataset: str,
217 data_dir: Path,
218 seed: int,
219 fractions: tuple[float, float, float],
220 preview: int,
221 package_version: str,
222) -> dict[str, Any]:
223 from tdc.generation import MolGen # Lazy optional import.
225 loader = MolGen(name=dataset, path=str(data_dir))
226 split = loader.get_split(method="random", seed=seed, frac=list(fractions))
227 expected = {"train", "valid", "test"}
228 if not isinstance(split, dict) or not expected.issubset(split):
229 raise CliError("PyTDC returned an unexpected MolGen split structure")
231 partitions: dict[str, Any] = {}
232 for name in ("train", "valid", "test"):
233 frame = split[name]
234 summary: dict[str, Any] = {
235 "columns": [str(column) for column in list(frame.columns)[:100]],
236 "rows": len(frame),
237 }
238 if preview:
239 summary["preview"] = truncate_value(
240 frame.head(preview).to_dict(orient="records")
241 )
242 partitions[name] = summary
244 return {
245 "action": "executed",
246 "data_directory": str(data_dir),
247 "dataset": dataset,
248 "download_acknowledged": True,
249 "package": "PyTDC",
250 "package_version": package_version,
251 "partitions": partitions,
252 "split": {
253 "fractions": list(fractions),
254 "method": "random",
255 "seed": seed,
256 },
257 }
260def _add_output_arguments(parser: argparse.ArgumentParser) -> None:
261 parser.add_argument("--output", help="write JSON to a relative workspace path")
262 parser.add_argument(
263 "--force", action="store_true", help="replace an existing --output file"
264 )
267def build_parser() -> argparse.ArgumentParser:
268 parser = argparse.ArgumentParser(
269 description=(
270 "Plan or explicitly execute bounded PyTDC molecular workflows. This "
271 "helper does not generate molecules and never calls remote-service or "
272 "docking oracles."
273 )
274 )
275 subparsers = parser.add_subparsers(dest="command", required=True)
277 score = subparsers.add_parser(
278 "score",
279 help="plan or run bounded scalar-oracle scoring",
280 description=(
281 "Score bounded user-supplied SMILES. QED is local. LogP, SA, and four "
282 "model-backed oracles require --download because they can fetch an "
283 "artifact/checkpoint."
284 ),
285 )
286 score.add_argument("--oracle", default="QED", help="exact package oracle name")
287 score.add_argument(
288 "--smiles", action="append", help="SMILES string; repeat for multiple inputs"
289 )
290 score.add_argument(
291 "--input", help="relative UTF-8 file with one SMILES string per line"
292 )
293 score.add_argument(
294 "--max-molecules",
295 type=bounded_int(1, 500),
296 default=100,
297 help="maximum accepted molecules (default: 100; max: 500)",
298 )
299 score.add_argument(
300 "--runtime-dir",
301 default=".pytdc-oracles",
302 help="relative directory for any acknowledged checkpoint (default: .pytdc-oracles)",
303 )
304 score.add_argument(
305 "--execute", action="store_true", help="execute the bounded scoring call"
306 )
307 score.add_argument(
308 "--download",
309 action="store_true",
310 help="acknowledge a supported model-checkpoint download",
311 )
312 _add_output_arguments(score)
314 dataset = subparsers.add_parser(
315 "dataset",
316 help="plan or load/split a MolGen corpus",
317 description=(
318 "Plan a MolGen random split. Execution requires both --execute and "
319 "--download because these corpora can be large."
320 ),
321 )
322 dataset.add_argument("--dataset", required=True, help="exact MolGen registry name")
323 dataset.add_argument(
324 "--frac",
325 nargs=3,
326 type=float,
327 metavar=("TRAIN", "VALID", "TEST"),
328 default=(0.7, 0.1, 0.2),
329 help="split fractions summing to 1 (default: 0.7 0.1 0.2)",
330 )
331 dataset.add_argument(
332 "--seed",
333 type=bounded_int(0, 4_294_967_295),
334 default=42,
335 help="non-negative random split seed (default: 42)",
336 )
337 dataset.add_argument(
338 "--data-dir",
339 default=".pytdc-molgen",
340 help="relative dataset cache directory (default: .pytdc-molgen)",
341 )
342 dataset.add_argument(
343 "--preview",
344 type=bounded_int(0, 5),
345 default=0,
346 help="include at most this many bounded rows per partition (default: 0)",
347 )
348 dataset.add_argument(
349 "--execute", action="store_true", help="construct and split the MolGen loader"
350 )
351 dataset.add_argument(
352 "--download",
353 action="store_true",
354 help="acknowledge potentially large dataset download/storage",
355 )
356 _add_output_arguments(dataset)
357 return parser
360def main(argv: list[str] | None = None) -> int:
361 args = build_parser().parse_args(argv)
362 try:
363 metadata, package_version = load_pytdc_metadata()
364 if args.command == "score":
365 smiles = load_smiles(
366 args.smiles, args.input, max_molecules=args.max_molecules
367 )
368 oracle, category = classify_oracle(metadata, args.oracle)
369 runtime_dir = safe_directory(
370 args.runtime_dir,
371 label="oracle runtime directory",
372 create=args.execute,
373 )
374 if args.execute:
375 result = execute_scores(
376 oracle_name=oracle,
377 category=category,
378 smiles=smiles,
379 runtime_dir=runtime_dir,
380 download_acknowledged=args.download,
381 package_version=package_version,
382 )
383 else:
384 result = score_plan(
385 oracle=oracle,
386 category=category,
387 smiles_count=len(smiles),
388 runtime_dir=runtime_dir,
389 package_version=package_version,
390 download_acknowledged=args.download,
391 )
392 else:
393 fractions = validate_fractions(args.frac)
394 registry_key = canonical_name("MolGen", metadata.dataset_names, "task")
395 dataset = canonical_name(
396 args.dataset, metadata.dataset_names[registry_key], "MolGen dataset"
397 )
398 data_dir = safe_directory(
399 args.data_dir,
400 label="MolGen data directory",
401 create=args.execute,
402 )
403 if args.execute:
404 if not args.download:
405 raise CliError(
406 "MolGen execution may download a large corpus; pass "
407 "--download together with --execute to acknowledge it"
408 )
409 result = execute_dataset(
410 dataset=dataset,
411 data_dir=data_dir,
412 seed=args.seed,
413 fractions=fractions,
414 preview=args.preview,
415 package_version=package_version,
416 )
417 else:
418 result = dataset_plan(
419 dataset=dataset,
420 data_dir=data_dir,
421 seed=args.seed,
422 fractions=fractions,
423 package_version=package_version,
424 download_acknowledged=args.download,
425 )
426 emit_json(result, args.output, force=args.force)
427 except (CliError, ImportError, OSError, TypeError, ValueError) as exc:
428 print(f"error: {exc}", file=sys.stderr)
429 return 2
430 return 0
433if __name__ == "__main__":
434 raise SystemExit(main())