From 42fba8451389bd7718cb514e53ad1feb5550c8a6 Mon Sep 17 00:00:00 2001 From: isayev Date: Tue, 4 Aug 2026 10:06:54 -0400 Subject: [PATCH 1/2] refactor!: split the oversized utils modules, and cut the optimizer's syncs Three lanes that interleave on SPE.py and ANI2xt_no_rep.py, so they land together rather than as a commit that cannot import. ## File splits utils/file_ops.py (849 lines) and utils/chemistry.py (580) are deleted with no shims, becoming utils/{smi_io,sdf_io,reconciliation,output_guard,geometry, connectivity,molprops,atomic_io}.py plus top-level id_mapping.py, job_layout.py, clash_relief.py and embedding.py. output_guard had to come out of validation.py first: validation.py imports torch and the model tree, so without the extraction a .smi writer pulled all of it in. A subprocess test asserts the file-I/O modules load neither. Phases 1 and 2 are pure moves, and the suite is the oracle for that: exactly one assertion changed, in a test reading hartree2ev off the module that no longer exists. Loggers are preserved byte-for-byte, since chemistry.py used getLogger("auto3d") and a renamed logger would silently stop reaching the run log. isomers/ loses its adapter shell and, with parallel_embed moved out to embedding.py, its import cycle: isomer_engine.py now imports nothing under Auto3D.isomers at any scope, enforced by an AST test. Engine construction stays deferred to run() deliberately -- RDKitIsomer.__init__ calls rdk_tmp.mkdir(), so constructing eagerly would move a filesystem side effect to a new point in the run. Two tests assert construction touches no filesystem. test_isomers.py needed real assertion changes, not import churn: it asserted isinstance against the adapter classes being deleted. A spy fixture now pins the full kwarg mapping per engine, which is strictly stronger than the attribute reads it replaces. ## One way to replace a file utils/atomic_io.py's atomic_write_path stages a sibling temp, copies the target's mode onto it, and os.replace()s. reorder_sdf did neither: it used a predictable .reorder.tmp name and dropped the target's mode, so a 0600 input came back 0664 -- a permission loosening. Confirmed as a failing test first, then fixed, across three call sites. The test that was supposed to guard this scanned for ".tmp" in the filename, which mkstemp's tmpXXXXXXXX.sdf never matches, so it would have passed either way. Rewritten as an exclusion scan. ## Cleanup ev2hatree was defined twice, misspelled, and is now EV_TO_HARTREE in constants.py. smiles2mols built torch.device(f"cuda:{idx}") by hand with no bounds check, so an out-of-range index produced cuda:99 instead of an error; it now goes through get_device, which raises GPUError. main()'s docstring advertised a SystemExit it never raises. The name-or-default idiom repeated 11 times in thermo.py is one helper. Three duplicated CLI error paths collapse. ## Optimizer syncs The FIRE loop did 18 host-device syncs per step from boolean-mask indexing. One nonzero() now feeds index_select reads and index_copy_ writes, and the smallest_fmax/oscillating_count block becomes two torch.where calls needing no index at all -- torch.where and not torch.minimum, because `<` is False for NaN and minimum would change NaN semantics. 18 syncs become 2, not 0: nonzero is itself a sync, and clean() needs a second one. Bit-identity is proven against a test-local reimplementation of the old loop running in the same process -- 17 scenarios x 5 state tensors, all torch.equal -- with a meta-test asserting the staggered scenario really does shrink the active set in stages, so the gathers are genuinely partial. One deliberate exception to bit-identity, and it fixes a pre-existing crash: the old loop assigned into smallest_fmax with a bool mask, which raises on a custom NNP returning float64 forces, but only with two or more molecules reducing at once -- a single one took the masked_fill_ fast path and cast silently. The explicit .to(dtype) at each index_copy_ makes that work. clean() now takes an int64 index and rejects a bool mask, which would otherwise have been reinterpreted as the indices [1,1,0,0]. n_steps, 200 lines on the hot path, decomposes into four helpers. Step-for-step bit-identity over 40 prefix lengths x 2 padding modes x 5 tensors. The ANI2xt element loop broke torch.compile into zero graphs, not seven: the data-dependent branch sits inside the loop, so Dynamo skipped the whole frame. It now compiles to one graph and satisfies fullgraph=True. No speedup is claimed anywhere. Every count here is exact and CI-enforced on CPU; no duration is measurable on a box whose GPUs are all busy. benchmarks/run_perf_ab.sh is one command for a maintainer with a free GPU, and it aborts rather than reporting if the two trees are the same, the hardware differs, or the converged counts move. Four tracked files claimed a ~1.25x torch.compile speedup for a path that compiled to zero graphs. They now say it is off by default, that no figure has been measured, and how to measure one. Verified: 1576 passed, 9 skipped, 1 xfailed, ruff clean, three randomized orderings. The remaining xfail is the E_tot filter divergence, which the next phase closes. --- CHANGELOG.md | 104 ++ benchmarks/.gitignore | 1 + benchmarks/bench_optimization_perf.py | 446 ++++++ benchmarks/run_perf_ab.sh | 48 + docs/source/advanced_usage.rst | 16 +- docs/source/howto/hpc.rst | 8 +- docs/source/migration.rst | 4 +- src/Auto3D/ASE/geometry.py | 87 +- src/Auto3D/ASE/thermo.py | 50 +- src/Auto3D/SPE.py | 6 +- src/Auto3D/auto3D.py | 50 +- src/Auto3D/batch_opt/ANI2xt_no_rep.py | 205 ++- src/Auto3D/batch_opt/batchopt.py | 4 +- src/Auto3D/batch_opt/fire_optimizer.py | 40 +- src/Auto3D/batch_opt/optimization_engine.py | 444 ++++-- src/Auto3D/chunk_manager.py | 2 +- src/Auto3D/clash_relief.py | 105 ++ src/Auto3D/cli/commands/config.py | 31 +- src/Auto3D/cli/commands/models.py | 4 +- src/Auto3D/cli/commands/run.py | 18 +- src/Auto3D/config.py | 2 +- src/Auto3D/constants.py | 6 + .../parallel_embed.py => embedding.py} | 17 +- src/Auto3D/filtering.py | 113 +- src/Auto3D/id_mapping.py | 180 +++ src/Auto3D/isomer_engine.py | 13 +- src/Auto3D/isomers/__init__.py | 34 +- src/Auto3D/isomers/base.py | 41 +- src/Auto3D/isomers/factory.py | 175 ++- src/Auto3D/isomers/omega_adapter.py | 71 - src/Auto3D/isomers/rdkit_adapters.py | 140 -- src/Auto3D/job_layout.py | 126 ++ src/Auto3D/models/adapter.py | 72 +- src/Auto3D/processors.py | 4 +- src/Auto3D/ranking.py | 15 +- src/Auto3D/tautomer.py | 5 +- src/Auto3D/utils/__init__.py | 35 +- src/Auto3D/utils/atomic_io.py | 101 ++ src/Auto3D/utils/chemistry.py | 580 -------- src/Auto3D/utils/connectivity.py | 207 +++ src/Auto3D/utils/convergence.py | 2 +- src/Auto3D/utils/energy.py | 24 +- src/Auto3D/utils/file_ops.py | 849 ------------ src/Auto3D/utils/geometry.py | 111 ++ src/Auto3D/utils/logging_config.py | 2 +- src/Auto3D/utils/molprops.py | 75 + src/Auto3D/utils/output_guard.py | 155 +++ src/Auto3D/utils/reconciliation.py | 170 +++ src/Auto3D/utils/sdf_io.py | 206 +++ src/Auto3D/utils/smi_io.py | 238 ++++ src/Auto3D/utils/stereochemistry.py | 55 +- src/Auto3D/utils/validation.py | 145 +- src/Auto3D/workflow.py | 45 +- src/Auto3D/workflow_workers.py | 4 +- tests/helpers_pipeline_output.py | 6 +- tests/helpers_sync_count.py | 223 +++ tests/test_ani2xt_atom_energies.py | 376 +++++ tests/test_atomic_io.py | 261 ++++ tests/test_config.py | 2 +- tests/test_config_parity.py | 2 +- tests/test_durability.py | 86 +- tests/test_e_tot_units.py | 2 +- tests/test_filter_unique.py | 268 ++++ tests/test_filtering.py | 8 +- tests/test_fire_optimizer.py | 60 +- tests/test_id_mapping.py | 197 +++ tests/test_import_boundaries.py | 99 +- tests/test_isomer_engine.py | 3 +- tests/test_isomer_engine_hardening.py | 7 +- tests/test_isomers.py | 394 +++--- tests/test_job_layout.py | 150 ++ tests/test_optimization_engine.py | 10 +- tests/test_optimization_engine_indexing.py | 834 +++++++++++ tests/test_parallel_embed.py | 8 +- tests/test_pipeline_e2e.py | 2 +- tests/test_processors.py | 24 +- tests/test_ranking.py | 6 +- tests/test_stereo_identity.py | 4 +- tests/test_stereo_postopt.py | 2 +- tests/test_thermo_imaginary_mode_inversion.py | 7 +- tests/test_thermo_reference.py | 75 +- tests/test_utils.py | 4 +- tests/test_utils_chemistry.py | 723 ---------- tests/test_utils_connectivity.py | 211 +++ tests/test_utils_energy_constants.py | 66 + tests/test_utils_file_ops.py | 1223 ----------------- tests/test_utils_geometry.py | 155 +++ ...st_chemistry.py => test_utils_molprops.py} | 48 +- tests/test_utils_reconciliation.py | 255 ++++ tests/test_utils_sdf_io.py | 321 +++++ tests/test_utils_smi_io.py | 363 +++++ tests/test_utils_validation.py | 3 +- tests/test_validate_run_parity.py | 4 +- tests/test_workflow.py | 16 +- 94 files changed, 7588 insertions(+), 4611 deletions(-) create mode 100644 benchmarks/.gitignore create mode 100644 benchmarks/bench_optimization_perf.py create mode 100755 benchmarks/run_perf_ab.sh create mode 100644 src/Auto3D/clash_relief.py rename src/Auto3D/{isomers/parallel_embed.py => embedding.py} (89%) create mode 100644 src/Auto3D/id_mapping.py delete mode 100644 src/Auto3D/isomers/omega_adapter.py delete mode 100644 src/Auto3D/isomers/rdkit_adapters.py create mode 100644 src/Auto3D/job_layout.py create mode 100644 src/Auto3D/utils/atomic_io.py delete mode 100644 src/Auto3D/utils/chemistry.py create mode 100644 src/Auto3D/utils/connectivity.py delete mode 100644 src/Auto3D/utils/file_ops.py create mode 100644 src/Auto3D/utils/geometry.py create mode 100644 src/Auto3D/utils/molprops.py create mode 100644 src/Auto3D/utils/output_guard.py create mode 100644 src/Auto3D/utils/reconciliation.py create mode 100644 src/Auto3D/utils/sdf_io.py create mode 100644 src/Auto3D/utils/smi_io.py create mode 100644 tests/helpers_sync_count.py create mode 100644 tests/test_ani2xt_atom_energies.py create mode 100644 tests/test_atomic_io.py create mode 100644 tests/test_filter_unique.py create mode 100644 tests/test_id_mapping.py create mode 100644 tests/test_job_layout.py create mode 100644 tests/test_optimization_engine_indexing.py delete mode 100644 tests/test_utils_chemistry.py create mode 100644 tests/test_utils_connectivity.py create mode 100644 tests/test_utils_energy_constants.py delete mode 100644 tests/test_utils_file_ops.py create mode 100644 tests/test_utils_geometry.py rename tests/{test_chemistry.py => test_utils_molprops.py} (65%) create mode 100644 tests/test_utils_reconciliation.py create mode 100644 tests/test_utils_sdf_io.py create mode 100644 tests/test_utils_smi_io.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 643af971..004b3b58 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -826,8 +826,112 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 use `Auto3D.ranking.species_id`. Runs made with `enumerate_isomer=False` and two inputs sharing an InChIKey lost a molecule and should be re-run. +### Performance + +- **The optimization loop no longer subsets the batch with boolean masks: 18 + host-device synchronizations per step become 2.** `n_steps` gathered the + still-active structures with `state['coord'][not_converged]` and friends -- + six masked reads, six masked writes, and four more inside `FIRE.clean`. On + CUDA every one of those is a GPU->CPU synchronization, because ATen has to + `nonzero()` the mask and copy the match count to the host to size the output. + Measured with a dispatch-mode counter: exactly 18 per step, on every step of + a loop that runs up to 2000 steps per bucket. + + The loop now computes `torch.nonzero(not_converged)` **once** and feeds the + resulting int64 index to `index_select` (reads) and `index_copy_` (writes), + neither of which synchronizes. The `smallest_fmax` and oscillation-counter + updates need no index at all and became `torch.where`. + + **This is 18 -> 2, not 18 -> 0.** `nonzero` *is* the synchronization -- it is + the mechanism by which boolean-mask indexing synced in the first place. The + win is that one `nonzero` result is reused by twelve gathers and scatters + instead of each computing its own, plus a second `nonzero` for `FIRE.clean`, + whose mask is indexed within the active subset rather than the full batch. + + **Results are bit-identical.** `index_select(0, nonzero(m))` and `x[m]` gather + the same rows in the same order (`nonzero` returns ascending indices), and + `index_copy_` writes the same rows as `x[m] = v`. No arithmetic, dtype or + reduction order changed. `tests/test_optimization_engine_indexing.py` runs a + test-local reimplementation of the old boolean-mask loop against the new one in + the same process and asserts `torch.equal` on coordinates, energies, fmax, + convergence mask and oscillation counters across 17 scenarios -- staggered + convergence, oscillation drops, a single molecule, `n=0`, batch 64, a padded + batch, and ten random seeds. `torch.where` rather than `torch.minimum` for + `smallest_fmax`, because `<` is False for NaN and the masked assignment it + replaced therefore *kept* the previous value; `minimum` would propagate the + NaN. + + **No speedup is claimed here.** The sync count is a fact this repository can + prove and CI enforces; the wall-clock value of removing a synchronization is + not, because it depends on the ratio of CPU launch time to GPU work at a given + batch size. `benchmarks/run_perf_ab.sh ` is one command that + measures it on a real GPU, sweeping batch 8/64/256/1024 across three molecule + sizes, and prints a block for this file. It aborts rather than reporting a + ratio if the two sides converge differently. + +- **ANI2xt's per-element energy loop: 22 host-device synchronizations per + forward become 2, and it compiles for the first time.** `ANI2xt.forward` + looped over its seven networks doing + `if mask.any(): atom_energies[mask] = network(aev[mask])` -- 7 guard + readbacks, 7 masked reads and 7 masked writes, plus 1 in `_validate_outputs`. + The guard protected nothing: `network(empty)` returns an empty tensor and the + write is a no-op, which is why deleting it is bit-identical even for a batch + containing 2 of the 7 elements. + + The loop is now `index_select`/`index_copy` over a flattened atom axis with + per-element indices computed by `element_indices()`, which reproduces seven + `nonzero` calls exactly (same indices, same order, verified over 200+ species + patterns including padding and out-of-range values) using a single host + readback of a fixed-size count vector. Self-atomic energies, a pure function + of species that was recomputed on every forward, moved to + `self_atomic_energies()`. + + **`compile_model=True` / `AUTO3D_COMPILE_MODEL=1` compiled *zero* subgraphs + for this model.** `if mask.any():` is a data-dependent branch, and a graph + break inside a `for` loop gives Dynamo nowhere to place a resume point, so it + skipped the entire frame. Deleting the guard alone does not fix that -- + `nonzero` and boolean-mask indexing are dynamic-output-shape ops and break the + same way -- which is why the indices are computed outside `forward` and passed + in. The per-element loop now compiles to one subgraph and passes + `fullgraph=True` (`tests/test_ani2xt_atom_energies.py`, which needs neither a + GPU nor torchani). Whether the *whole* `forward`, with torchani's AEVComputer + in the frame, also reaches one subgraph is a torchani-only measurement that + `benchmarks/bench_optimization_perf.py` reports. + +- **A custom NNP returning float64 forces no longer crashes the optimizer.** + `smallest_fmax` is allocated float32 and `fmax` inherited float64 from the + forces, so `smallest_fmax[reduced] = fmax[reduced]` raised `"Index put + requires the source and destination dtypes match"` -- but only when two or + more structures reduced their force in the same step, since a single-element + value took ATen's `masked_fill_` fast path and silently cast. The failure was + therefore batch-size dependent and invisible to a single-molecule test. Every + state write now casts explicitly to the destination dtype. + +### Changed + +- **The documented "~1.25x" `torch.compile` speedup for ANI models has been + removed from the docs, because no measurement supports it.** It appeared in + `docs/source/advanced_usage.rst`, `docs/source/migration.rst`, + `docs/source/howto/hpc.rst` and two adapter docstrings. For ANI2xt it cannot + have originated in the model at all: as described above, that path compiled + zero subgraphs. The docs now state that the setting is off by default, that no + figure has been measured, and how to measure one. + ### Added +- **`benchmarks/bench_optimization_perf.py` and `benchmarks/run_perf_ab.sh`.** + One command -- `bash benchmarks/run_perf_ab.sh v4.0.0` -- creates a read-only + git worktree of the base ref, benchmarks it and the current tree on the same + GPU with identical instrumentation, and prints a CHANGELOG-ready block. Fixed + work (`opttol=0`, `patience=1e9`) so both sides execute the same number of + full-width steps; 3x20-step warmup discarded; 7 reps reported as median with + IQR; rows whose IQR exceeds 10% of the median are flagged noisy and excluded + from the summary; sync counts come from a separate `set_sync_debug_mode` pass + because the instrumentation perturbs timing. It aborts if both runs import + Auto3D from the same tree, if the hardware differs, if there is no GPU, or if + converged counts or energies moved -- and the summary quotes a range across + batch sizes rather than a best case. + - **`auto3d validate` accepts `--json`**, the one result-producing command that did not have it while `run`, `energy`, `optimize`, `thermo` and `tautomers` all did. The document reports `success`, `format`, `molecules`, diff --git a/benchmarks/.gitignore b/benchmarks/.gitignore new file mode 100644 index 00000000..fbca2253 --- /dev/null +++ b/benchmarks/.gitignore @@ -0,0 +1 @@ +results/ diff --git a/benchmarks/bench_optimization_perf.py b/benchmarks/bench_optimization_perf.py new file mode 100644 index 00000000..85dec3f6 --- /dev/null +++ b/benchmarks/bench_optimization_perf.py @@ -0,0 +1,446 @@ +#!/usr/bin/env python +"""Measure the wall-clock and host-device-sync cost of one Auto3D optimization step. + +This is the *only* thing in this repository that produces a speedup number, and +it must be run on a real GPU with real neural network potentials. Nothing in the +test suite can do it: CI has no GPU, and a host-device synchronization costs only +what it serializes, so its price depends entirely on the ratio of CPU launch time +to GPU work at a given batch size. The tests count syncs exactly; this prices +them. + +Usage -- one command does everything:: + + bash benchmarks/run_perf_ab.sh v4.0.0 + +or, if you would rather git not be touched by a script:: + + git checkout v4.0.0 && python benchmarks/bench_optimization_perf.py --label before + git checkout && python benchmarks/bench_optimization_perf.py --label after + python benchmarks/bench_optimization_perf.py --compare before after + +``--compare`` prints a block meant to be pasted into the CHANGELOG verbatim. It +refuses to print anything if the two runs imported the same source tree, or if +the optimization outcomes differ between them. + +Design notes, because each of these is the difference between a number and a +number you can trust: + +* **Fixed work.** The timed loop uses ``opttol=0.0`` and ``patience=10**9`` so no + structure ever leaves the active set and every configuration executes exactly + ``--steps`` full-width steps. Comparing runs that converge at different steps + measures convergence luck, not throughput. +* **Batch sweep 8/64/256/1024.** Syncs dominate in the small-batch, + CPU-launch-bound regime and amortize when each step has enough GPU work. + Reporting only the batch size that looks best would be cherry-picking, so every + row is printed and the summary quotes a *range*. +* **Separate sync pass.** ``set_sync_debug_mode`` perturbs timing, so it never + shares a pass with the clock. +* **Warmup then 7 reps, median and IQR.** A row whose IQR exceeds 10% of its + median is flagged noisy and excluded from the summary. +* **Outcome gate.** A realism pass runs production settings and records converged + counts and energies. If those move, ``--compare`` aborts rather than reporting + a speedup, because a faster loop that computes something else is not faster. +""" +from __future__ import annotations + +import argparse +import json +import platform +import statistics +import subprocess +import sys +import time +import warnings +from collections import defaultdict +from pathlib import Path + +import torch + +RESULTS = Path(__file__).parent / "results" + +# 24 fixed drug-like SMILES: 8 small (~12-18 atoms), 8 medium (~30-40), 8 large +# (~55-70). Embedded rather than read from a file so the benchmark is identical +# across checkouts and needs no network. +SMILES = { + "small": ["CCO", "c1ccccc1O", "CC(=O)N", "CC(N)C(=O)O", "c1ccncc1", + "CSCC", "OCC(O)CO", "Clc1ccccc1"], + "medium": ["CC(=O)Nc1ccc(O)cc1", "CN1CCC[C@H]1c1cccnc1", + "c1ccc2[nH]ccc2c1C(=O)NCC", "OC(=O)c1ccccc1OC(C)=O", + "CC(C)Cc1ccc(cc1)C(C)C(=O)O", "Fc1ccc(cc1)C(=O)CCCN1CCCCC1", + "CCN(CC)CCNC(=O)c1ccc(N)cc1", "CC(C)(C)NCC(O)c1ccc(O)c(CO)c1"], + "large": ["CC1(C)SC2C(NC(=O)Cc3ccccc3)C(=O)N2C1C(=O)O", + "CN1C2CCC1C(C(=O)OC)C(OC(=O)c1ccccc1)C2", + "Cc1ccc(cc1)S(=O)(=O)NC(=O)NN1CCCCCC1", + "CC(=O)OC1CC2CCC3C(CCC4(C)C3CCC24C)C1", + "COc1cc2c(cc1OC)C(=O)C(CC2)Cc1ccc(OC)cc1", + "OC(=O)C1CCCN1C(=O)C(Cc1ccccc1)NC(=O)OCc1ccccc1", + "CN(C)CCCN1c2ccccc2CCc2ccccc21", + "CC(C)NCC(O)COc1cccc2ccccc12"], +} +BATCHES = (8, 64, 256, 1024) +N_STEPS_TIMED = 200 +N_WARMUP_CALLS, N_WARMUP_STEPS = 3, 20 +N_REPS = 7 +SEED = 0xA173D +NOISE_FRACTION = 0.10 +ENERGY_TOLERANCE_EV = 1e-4 + + +def build_mols() -> dict[str, list]: + """Deterministic 3D conformers via RDKit ETKDGv3. No network, no files.""" + from rdkit import Chem + from rdkit.Chem import AllChem + + out: dict[str, list] = {} + for size, smis in SMILES.items(): + mols = [] + for smi in smis: + mol = Chem.AddHs(Chem.MolFromSmiles(smi)) + params = AllChem.ETKDGv3() + params.randomSeed = SEED + if AllChem.EmbedMolecule(mol, params) != 0: + raise SystemExit(f"embedding failed for {smi!r}; fix the SMILES set") + mol.SetProp("_Name", smi) + mols.append(mol) + out[size] = mols + return out + + +def env_block() -> dict: + """Capture everything needed to tell two runs apart, and to reproduce one.""" + def git(*args: str) -> str: + try: + return subprocess.check_output(["git", *args], text=True, + stderr=subprocess.DEVNULL).strip() + except Exception: + return "unknown" + + import Auto3D + + on_gpu = torch.cuda.is_available() + return { + "gpu": torch.cuda.get_device_name(0) if on_gpu else "CPU-ONLY", + "torch": torch.__version__, + "cuda": torch.version.cuda, + "capability": (".".join(map(str, torch.cuda.get_device_capability(0))) + if on_gpu else "n/a"), + "python": platform.python_version(), + "commit": git("rev-parse", "--short", "HEAD"), + "ref": git("rev-parse", "--abbrev-ref", "HEAD"), + # The A/B guard: if these match, both runs benchmarked the same code. + "auto3d_path": str(Path(Auto3D.__file__).resolve().parent), + } + + +def make_state(mols: list, batch: int, device: torch.device, model) -> tuple[dict, torch.Tensor, int]: + """Pad one bucket of ``batch`` molecules (cycled) into an ``n_steps`` state.""" + from Auto3D.batch_opt.model_wrapper import EnForce_ANI + from Auto3D.batch_opt.padding import pad_from_mols + from Auto3D.constants import INITIAL_ENERGY_SENTINEL, INITIAL_FMAX_SENTINEL + + picked = [mols[i % len(mols)] for i in range(batch)] + coord, numbers, charges, atom_mask = pad_from_mols(picked, model, device) + coord = coord.detach().to(dtype=torch.float, device=device) + size = coord.shape[0] + state = { + "coord": coord, + "numbers": numbers, + "charges": charges, + "nn": EnForce_ANI(model, 1024 * 16), + "converged_mask": torch.zeros(size, dtype=torch.bool, device=device), + "fmax": torch.full((size,), INITIAL_FMAX_SENTINEL, device=device), + "energy": torch.full((size,), INITIAL_ENERGY_SENTINEL, dtype=torch.double, + device=device), + "timing": defaultdict(float), + } + return state, atom_mask, int(coord.shape[1]) + + +def _sync() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def time_steps(mols, batch, device, model, steps, reps) -> dict: + """Median/min microseconds per step over ``reps``, with nothing converging.""" + from Auto3D.batch_opt.optimization_engine import n_steps + + for _ in range(N_WARMUP_CALLS): + state, mask, _ = make_state(mols, batch, device, model) + n_steps(state, n=N_WARMUP_STEPS, opttol=0.0, patience=10 ** 9, atom_mask=mask) + _sync() + + per_step = [] + natoms = 0 + for _ in range(reps): + state, mask, natoms = make_state(mols, batch, device, model) + _sync() + start = time.perf_counter() + # Brackets only the outside of the rep. Synchronizing inside would add + # exactly the syncs under study. + n_steps(state, n=steps, opttol=0.0, patience=10 ** 9, atom_mask=mask) + _sync() + per_step.append((time.perf_counter() - start) / steps * 1e6) + + quartiles = statistics.quantiles(per_step, n=4) if len(per_step) >= 4 else [0, 0, 0] + return {"median_us": statistics.median(per_step), "min_us": min(per_step), + "iqr_us": quartiles[2] - quartiles[0], "natoms": natoms, + "reps": per_step} + + +def count_syncs(mols, batch, device, model, steps=10) -> float | None: + """Per-step sync count from ``set_sync_debug_mode('warn')``. + + A separate pass from timing: the instrumentation perturbs what it measures. + Returns None without CUDA, where the concept does not apply -- the CPU-side + equivalent is counted exactly by ``tests/test_optimization_engine_indexing.py``. + """ + from Auto3D.batch_opt.optimization_engine import n_steps + + if not torch.cuda.is_available(): + return None + + state, mask, _ = make_state(mols, batch, device, model) + n_steps(state, n=2, opttol=0.0, patience=10 ** 9, atom_mask=mask) # warm + state, mask, _ = make_state(mols, batch, device, model) + _sync() + torch.cuda.set_sync_debug_mode("warn") + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + n_steps(state, n=steps, opttol=0.0, patience=10 ** 9, atom_mask=mask) + hits = [w for w in caught if "synchron" in str(w.message).lower()] + return len(hits) / steps + finally: + torch.cuda.set_sync_debug_mode("default") + + +def count_subgraphs(device) -> int | str | None: + """Compiled subgraphs in ``ANI2xt.forward``. None if torchani is absent. + + This is the claim the test suite cannot close: ``tests/test_ani2xt_atom_energies.py`` + proves the per-element loop compiles to one subgraph in isolation, but with + torchani's AEVComputer in the same frame the count could still be zero, since + any break inside that ``for`` loop makes Dynamo skip the whole frame. + """ + try: + import torchani # noqa: F401 + except ImportError: + return None + + import torch._dynamo as dynamo + + from Auto3D.batch_opt.ANI2xt_no_rep import ANI2xt, element_indices, self_atomic_energies + + graphs: list = [] + + def backend(gm, example_inputs): + graphs.append(gm) + return gm.forward + + model = ANI2xt(device) + dynamo.reset() + species = torch.zeros(2, 8, dtype=torch.long, device=device) + species[0, :4] = 1 + coords = torch.randn(2, 8, 3, device=device, requires_grad=True) + kwargs = {} + try: + kwargs = { + "elem_index": element_indices(species, len(model.networks)), + "self_energies": self_atomic_energies(species, model.energy_shifts, + len(model.networks)), + } + except Exception: + pass + try: + torch.compile(model, backend=backend, dynamic=True, + fullgraph=False)(species, coords, **kwargs) + except Exception as exc: + return f"error: {type(exc).__name__}: {exc}" + finally: + dynamo.reset() + return len(graphs) + + +def realism_pass(mols, device, model) -> dict: + """Production settings end to end, so ``--compare`` can verify outcomes match.""" + from Auto3D.batch_opt.optimization_engine import n_steps + + everything = [m for group in mols.values() for m in group] + state, mask, _ = make_state(everything, len(everything), device, model) + _sync() + start = time.perf_counter() + n_steps(state, n=2000, opttol=0.01, patience=250, atom_mask=mask) + _sync() + return {"wall_s": time.perf_counter() - start, + "converged": int(state["converged_mask"].sum().item()), + "total": len(everything), + "energies": [float(x) for x in state["energy"].tolist()]} + + +def run(label: str, engines: list[str], device_str: str, steps: int, reps: int) -> None: + """Benchmark every engine x size x batch combination and write a JSON record.""" + from Auto3D.model_factory import create_model + + device = torch.device(device_str if torch.cuda.is_available() else "cpu") + mols = build_mols() + record = {"label": label, "env": env_block(), "steps": steps, "reps": reps, + "rows": [], "extra": {}} + + if not torch.cuda.is_available(): + print("WARNING: no CUDA device. Sync counts are unavailable and timings " + "are NOT representative of the change under test. Do not report " + "these numbers.", file=sys.stderr) + + for engine in engines: + compile_model = engine.endswith("+compile") + name = engine.removesuffix("+compile") + try: + model = create_model(name, device, compile_model=compile_model) + except Exception as exc: + print(f" SKIP {engine}: {type(exc).__name__}: {exc}", file=sys.stderr) + record["extra"][f"skipped.{engine}"] = f"{type(exc).__name__}: {exc}" + continue + + for size, group in mols.items(): + for batch in BATCHES: + timing = time_steps(group, batch, device, model, steps, reps) + syncs = count_syncs(group, batch, device, model) + record["rows"].append({"engine": engine, "size": size, + "batch": batch, "syncs_per_step": syncs, + **timing}) + print(f" {engine:16s} {size:6s} b={batch:<5d} " + f"{timing['median_us']:9.1f} us/step syncs/step={syncs}") + record["extra"][f"realism.{engine}"] = realism_pass(mols, device, model) + del model + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + record["extra"]["ani2xt_subgraphs"] = count_subgraphs(device) + RESULTS.mkdir(parents=True, exist_ok=True) + out = RESULTS / f"{label}.json" + out.write_text(json.dumps(record, indent=2)) + print(f"\nwrote {out}") + + +def _load(label: str) -> dict: + path = RESULTS / f"{label}.json" + if not path.exists(): + sys.exit(f"ABORT: {path} does not exist. Run --label {label} first.") + return json.loads(path.read_text()) + + +def compare(before_label: str, after_label: str) -> None: + """Print the CHANGELOG block, or abort. Never both.""" + before, after = _load(before_label), _load(after_label) + + for record in (before, after): + if "env" not in record or record["env"].get("commit") == "unknown": + sys.exit("ABORT: a result file has no usable environment block. " + "Re-run both sides inside a git checkout.") + if before["env"]["auto3d_path"] == after["env"]["auto3d_path"]: + sys.exit("ABORT: both runs imported Auto3D from " + f"{before['env']['auto3d_path']}. Nothing was compared. Use " + "benchmarks/run_perf_ab.sh, which points PYTHONPATH at two trees.") + if before["env"]["gpu"] != after["env"]["gpu"]: + sys.exit(f"ABORT: different hardware ({before['env']['gpu']} vs " + f"{after['env']['gpu']}). A cross-machine ratio is not a speedup.") + if before["env"]["gpu"] == "CPU-ONLY": + sys.exit("ABORT: these runs had no GPU. The change under test removes " + "host-device synchronizations, which do not exist on CPU. " + "Report nothing.") + + # Outcome equality gate. A faster loop that computes something else is not + # faster, so this refuses to print a ratio at all. + for key, before_extra in before["extra"].items(): + if not key.startswith("realism."): + continue + after_extra = after["extra"].get(key) + if after_extra is None: + continue + if before_extra["converged"] != after_extra["converged"]: + sys.exit(f"ABORT: OUTCOMES CHANGED for {key}: converged " + f"{before_extra['converged']} -> {after_extra['converged']}. " + "Do not report a speedup.") + worst = max(abs(x - y) for x, y in + zip(before_extra["energies"], after_extra["energies"], strict=True)) + if worst > ENERGY_TOLERANCE_EV: + sys.exit(f"ABORT: OUTCOMES CHANGED for {key}: max |dE| = {worst:.3e} eV " + f"> {ENERGY_TOLERANCE_EV:.0e}. Do not report a speedup.") + + env = after["env"] + print("### Performance (measured, not estimated)\n") + print(f"Host: {env['gpu']} (sm_{env['capability']}) | torch {env['torch']} " + f"| CUDA {env['cuda']} | python {env['python']}") + print(f"Before: {before['env']['ref']} ({before['env']['commit']}) " + f"After: {env['ref']} ({env['commit']})") + print(f"Fixed-work loop: {after['steps']} steps, opttol=0, patience=1e9, " + f"{after['reps']} reps, median of per-step wall clock.\n") + print("| engine | mol size | batch | atoms/mol | us/step before | us/step after " + "| speedup | syncs/step before | after |") + print("|---|---|---|---|---|---|---|---|---|") + + by_key = {(r["engine"], r["size"], r["batch"]): r for r in before["rows"]} + speedups = [] + for row in after["rows"]: + key = (row["engine"], row["size"], row["batch"]) + if key not in by_key: + continue + base = by_key[key] + ratio = base["median_us"] / row["median_us"] + noisy = (base["iqr_us"] > NOISE_FRACTION * base["median_us"] + or row["iqr_us"] > NOISE_FRACTION * row["median_us"]) + if not noisy: + speedups.append(ratio) + print(f"| {row['engine']} | {row['size']} | {row['batch']} | " + f"{row['natoms']} | {base['median_us']:.1f} | {row['median_us']:.1f} | " + f"{ratio:.2f}x{' (noisy)' if noisy else ''} | " + f"{base['syncs_per_step']} | {row['syncs_per_step']} |") + + print(f"\nCompiled subgraphs in ANI2xt.forward: " + f"before {before['extra'].get('ani2xt_subgraphs')}, " + f"after {after['extra'].get('ani2xt_subgraphs')}.") + + for key, before_extra in before["extra"].items(): + if key.startswith("realism.") and key in after["extra"]: + after_extra = after["extra"][key] + engine = key.split(".", 1)[1] + print(f"End-to-end {engine} (opttol=0.01, patience=250, " + f"{before_extra['total']} molecules): before " + f"{before_extra['wall_s']:.1f}s, after {after_extra['wall_s']:.1f}s " + f"({before_extra['wall_s'] / after_extra['wall_s']:.2f}x). " + f"Outcomes unchanged: converged {before_extra['converged']}/" + f"{before_extra['total']} -> {after_extra['converged']}/" + f"{after_extra['total']}.") + + if speedups: + print(f"\nSummary: {min(speedups):.2f}x-{max(speedups):.2f}x per optimization " + f"step across batch sizes {min(BATCHES)}-{max(BATCHES)} and three " + f"molecule sizes (noisy rows excluded). Quote the range, not the " + f"best row.") + else: + print("\nSummary: NO non-noisy rows. Re-run on an idle GPU. Report nothing.") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--label", help="Run the benchmark and save as this label.") + parser.add_argument("--engines", default="aimnet2,ANI2xt,ANI2xt+compile", + help="Comma-separated. Append '+compile' for torch.compile.") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--steps", type=int, default=N_STEPS_TIMED) + parser.add_argument("--reps", type=int, default=N_REPS) + parser.add_argument("--compare", nargs=2, metavar=("BEFORE", "AFTER")) + args = parser.parse_args() + + if args.compare: + compare(*args.compare) + elif args.label: + run(args.label, args.engines.split(","), args.device, args.steps, args.reps) + else: + parser.error("pass --label or --compare") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/run_perf_ab.sh b/benchmarks/run_perf_ab.sh new file mode 100755 index 00000000..ccd3b4a8 --- /dev/null +++ b/benchmarks/run_perf_ab.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +# benchmarks/run_perf_ab.sh +# +# ONE command: benchmark and the current working tree on the same GPU, +# then print a CHANGELOG-ready block. Nothing to interpret, nothing to assemble. +# +# bash benchmarks/run_perf_ab.sh v4.0.0 +# +# A read-only git worktree of is created in a temp dir and removed on +# exit. Auto3D is pure Python, so PYTHONPATH is enough -- no reinstall, and the +# two runs cannot contaminate each other's site-packages. +# +# Note that the *measurement code* comes from the current tree in both runs +# (only PYTHONPATH differs), which is deliberate: identical instrumentation, +# different measured code. bench_optimization_perf.py aborts if it detects both +# runs importing Auto3D from the same tree. +set -euo pipefail + +BASE="${1:?usage: run_perf_ab.sh (e.g. v4.0.0, or a SHA)}" +REPO="$(git rev-parse --show-toplevel)" +BENCH="$REPO/benchmarks/bench_optimization_perf.py" +TMP="$(mktemp -d)" +WT="$TMP/base" + +cleanup() { + git -C "$REPO" worktree remove --force "$WT" >/dev/null 2>&1 || true + rm -rf "$TMP" +} +trap cleanup EXIT + +if ! python -c 'import torch, sys; sys.exit(0 if torch.cuda.is_available() else 1)'; then + echo "ABORT: no CUDA device visible. This benchmark measures the removal of" >&2 + echo " host-device synchronizations, which do not exist on CPU." >&2 + exit 1 +fi + +echo "== creating read-only worktree of $BASE ==" +git -C "$REPO" worktree add --detach "$WT" "$BASE" >/dev/null + +echo "== baseline ($BASE) ==" +PYTHONPATH="$WT/src" python "$BENCH" --label before + +echo "== branch ($(git -C "$REPO" rev-parse --abbrev-ref HEAD)) ==" +PYTHONPATH="$REPO/src" python "$BENCH" --label after + +echo +echo "== comparison (paste the block below into CHANGELOG.md) ==" +PYTHONPATH="$REPO/src" python "$BENCH" --compare before after diff --git a/docs/source/advanced_usage.rst b/docs/source/advanced_usage.rst index 89a91bcf..a29dfedf 100644 --- a/docs/source/advanced_usage.rst +++ b/docs/source/advanced_usage.rst @@ -106,7 +106,7 @@ Python API: torch.compile() Optimization ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Enable PyTorch 2.0 compilation for ANI models (~1.25x speedup): +Enable PyTorch 2.0 compilation for ANI models (off by default): .. code:: console @@ -128,8 +128,18 @@ Python API: model = create_model("ANI2xt", device=device, compile_model=True) .. note:: - ``torch.compile()`` works best with ANI2x/ANI2xt. AIMNET uses optimized - JIT compilation internally. + ``torch.compile()`` applies only to ANI2x/ANI2xt. AIMNET uses its own + compilation internally and ignores this setting. + + No speedup figure is documented here because none has been measured on this + codebase. Earlier versions of these docs quoted "~1.25x"; that number had no + measurement behind it, and for ANI2xt it could not have come from the model + at all -- ``ANI2xt.forward`` used to contain a data-dependent branch + inside its per-element loop, which made Dynamo skip the entire frame and + compile **zero** subgraphs. That loop is now compilable (one subgraph), so a + gain is at least possible; whether there is one, and how large, is a GPU + measurement. Run ``benchmarks/run_perf_ab.sh`` to get a number for your + hardware -- it reports eager and compiled ANI2xt side by side. Batch Size Tuning ~~~~~~~~~~~~~~~~~ diff --git a/docs/source/howto/hpc.rst b/docs/source/howto/hpc.rst index 9f4971d0..f36f6c8b 100644 --- a/docs/source/howto/hpc.rst +++ b/docs/source/howto/hpc.rst @@ -234,7 +234,13 @@ For PyTorch 2.0+, enable compilation: export AUTO3D_COMPILE_MODEL=1 auto3d run input.smi --k=1 --gpu --engine=ANI2x -Provides ~1.25x speedup after warmup. +Off by default, and no speedup figure is documented because none has been +measured on this codebase. The "~1.25x" these docs used to quote was +unsubstantiated, and for ANI2xt it could not have originated in the model: +its per-element loop used to graph-break in a way that made Dynamo skip the +whole frame, compiling zero subgraphs. To measure it on your own hardware, run +``benchmarks/run_perf_ab.sh``, which times ANI2xt eager against ANI2xt +compiled. Engine Selection for Speed ~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/source/migration.rst b/docs/source/migration.rst index 77691c48..50289698 100644 --- a/docs/source/migration.rst +++ b/docs/source/migration.rst @@ -180,7 +180,9 @@ Environment variables New environment variables for runtime configuration: -- ``AUTO3D_COMPILE_MODEL=1`` - Enable torch.compile for ANI models (~1.25x speedup) +- ``AUTO3D_COMPILE_MODEL=1`` - Enable torch.compile for ANI2x/ANI2xt (off by + default; no speedup figure is documented because none has been measured -- + see :doc:`advanced_usage`) - ``AIMNET_CACHE_DIR`` - Override the AIMNet2 model download cache (default: ``~/.cache/aimnet``) .. note:: diff --git a/src/Auto3D/ASE/geometry.py b/src/Auto3D/ASE/geometry.py index 82162b4e..8ab14fbc 100644 --- a/src/Auto3D/ASE/geometry.py +++ b/src/Auto3D/ASE/geometry.py @@ -5,8 +5,6 @@ from __future__ import annotations import os -import stat -import tempfile from rdkit import Chem @@ -20,6 +18,7 @@ from Auto3D.model_factory import create_model, get_device from Auto3D.models.preflight import resolve_engine_name from Auto3D.torch_config import TorchConfig, configure_torch +from Auto3D.utils.atomic_io import atomic_write_path from Auto3D.utils.energy import E_TOT_HARTREE_PROP, E_TOT_PROP from Auto3D.utils.validation import ( check_engine_supports_molecules, @@ -31,44 +30,6 @@ __all__ = ["opt_geometry"] -def _stage_beside(target: str) -> str: - """Create an empty temp file in the same directory as ``target``. - - Same directory, because ``os.replace`` raises ``OSError`` across - filesystems. The temp file inherits ``target``'s permission bits so the - replaced file keeps the mode it had -- ``tempfile.mkstemp`` creates 0600, - which would otherwise silently tighten the user's output file. Setting the - mode before anything is written also preserves a read-only (0444) target's - protection, which ``rename(2)`` would otherwise bypass. - - The parent directory is resolved with ``realpath``, not ``abspath``: - ``abspath`` collapses ``..`` lexically, so a target like - ``/scratch/link/../out.sdf`` (where ``link`` points at another mount) - would stage the temp file in ``/scratch`` while the replace destination - really lives elsewhere -- and ``os.replace`` would fail with - ``EXDEV: Invalid cross-device link`` after a completed run. Only the - PARENT is resolved: ``os.replace`` acts on the final path component - itself, so following a symlinked ``target`` would pick the wrong - directory. - """ - directory = os.path.realpath(os.path.dirname(os.path.abspath(target))) - fd, tmp_path = tempfile.mkstemp(suffix=".sdf", dir=directory) - os.close(fd) - try: - os.chmod(tmp_path, stat.S_IMODE(os.stat(target).st_mode)) - except OSError: - # Best effort: a mode we cannot read is not a reason to abandon a - # completed optimization. Defensive rather than exercised -- the sole - # caller, `_annotate_and_rewrite`, only reaches here after - # `Chem.SDMolSupplier(target)` has already read the file, so `target` - # exists and is stat-able on every path that gets here today. An - # earlier version of this comment claimed the branch fires "in normal - # runs whenever `target` does not exist yet", which is not true of any - # current caller. - pass - return tmp_path - - def _annotate_and_rewrite(outpath: str) -> None: """Add the unit-labeled ``E_tot(Hartree)`` sibling in-place, atomically. @@ -89,9 +50,11 @@ def _annotate_and_rewrite(outpath: str) -> None: geometries to ``outpath``. Opening ``Chem.SDWriter(outpath)`` directly would truncate that file, so a failure partway through the rewrite would destroy a completed optimization run (C14). Stage into a sibling temp file - and ``os.replace`` it into position instead: ``os.replace`` is atomic on - POSIX and on Windows, so ``outpath`` is only ever the old complete file or - the new complete file, never a partial one. + and ``os.replace`` it into position instead -- which is what + :func:`Auto3D.utils.atomic_io.atomic_write_path` does, for this and the + other two in-place rewrites in Auto3D. ``os.replace`` is atomic on POSIX + and on Windows, so ``outpath`` is only ever the old complete file or the + new complete file, never a partial one. Staging does NOT by itself remove the Windows hazard from 74474ed, and an earlier version of this docstring wrongly claimed it did. ``reorder_sdf`` @@ -99,39 +62,29 @@ def _annotate_and_rewrite(outpath: str) -> None: was an open ``SDMolSupplier`` on the ``os.replace`` DESTINATION, which Windows refuses to overwrite (``PermissionError``/``WinError 5``) while a handle is held. This function reads ``outpath`` and then replaces it, so it - has the same exposure -- see the explicit release below. + has the same exposure -- see the explicit release below. Releasing the + handle stays the caller's duty; ``atomic_write_path`` cannot do it. """ supp = Chem.SDMolSupplier(outpath, removeHs=False) mols = list(supp) # Release the handle on `outpath` BEFORE os.replace targets it, exactly as - # utils/file_ops.py:743 does for reorder_sdf. Writing this as the anonymous + # utils/sdf_io.py does for reorder_sdf. Writing this as the anonymous # `list(Chem.SDMolSupplier(...))` would also work today -- the temporary's # refcount drops at the end of the statement -- but only under CPython's # refcounting, and it leaves the requirement invisible to the next person # who refactors this into a named variable. del supp - tmp_path = _stage_beside(outpath) - try: - with Chem.SDWriter(tmp_path) as f: - for mol in mols: - # Skip records that failed to re-parse or lack E_tot rather - # than crashing, which would discard the entire (already - # completed) optimization run on a single bad record. - if mol is None or not mol.HasProp(E_TOT_PROP): - continue - # Same number, stated in a name that carries its unit. No - # arithmetic: E_tot is already Hartree when it gets here. - mol.SetProp(E_TOT_HARTREE_PROP, mol.GetProp(E_TOT_PROP)) - f.write(mol) - os.replace(tmp_path, outpath) - except BaseException: - # BaseException, not Exception: a KeyboardInterrupt mid-write must not - # leave a stray .sdf beside the user's output. - try: - os.unlink(tmp_path) - except OSError: - pass - raise + with atomic_write_path(outpath, suffix=".sdf") as tmp_path, Chem.SDWriter(tmp_path) as f: + for mol in mols: + # Skip records that failed to re-parse or lack E_tot rather + # than crashing, which would discard the entire (already + # completed) optimization run on a single bad record. + if mol is None or not mol.HasProp(E_TOT_PROP): + continue + # Same number, stated in a name that carries its unit. No + # arithmetic: E_tot is already Hartree when it gets here. + mol.SetProp(E_TOT_HARTREE_PROP, mol.GetProp(E_TOT_PROP)) + f.write(mol) def opt_geometry( diff --git a/src/Auto3D/ASE/thermo.py b/src/Auto3D/ASE/thermo.py index ea3f0b6e..c535dbe6 100644 --- a/src/Auto3D/ASE/thermo.py +++ b/src/Auto3D/ASE/thermo.py @@ -28,6 +28,7 @@ DEFAULT_OPT_STEPS, DEFAULT_THERMO_CONVERGENCE_THRESHOLD, EV_PER_WAVENUMBER, + EV_TO_HARTREE, IMAGINARY_MODE_CUTOFF_CM, LINEARITY_MAX_PERP_ANGSTROM, LINEARITY_MOMENT_RATIO, @@ -39,7 +40,7 @@ from Auto3D.models.preflight import resolve_engine_name from Auto3D.models.species import to_model_species from Auto3D.torch_config import TorchConfig, configure_torch -from Auto3D.utils.chemistry import hartree2ev +from Auto3D.utils.energy import hartree2ev from Auto3D.utils.logging_config import get_logger from Auto3D.utils.validation import ( check_engine_supports_molecules, @@ -52,7 +53,6 @@ # TF32 settings are configured centrally via Auto3D.torch_config.configure_torch() # and the allow_tf32 option in Auto3DOptions. -ev2hatree = 1/hartree2ev #: SD property carrying the success/failure verdict for one thermo record. #: ``""`` means publishable; any non-empty value names the failure. This is the @@ -68,6 +68,21 @@ logger = get_logger(__name__) +def _mol_name(mol: Chem.Mol, default: str = "molecule") -> str: + """The molecule's ``_Name`` property, or ``default`` when it has none. + + Every diagnostic/warning site in this module needs a human-readable + identifier for a mol that may or may not carry ``_Name``, and used to + repeat ``mol.GetProp("_Name") if mol.HasProp("_Name") else `` + verbatim at each site (M64). The default itself is NOT hardcoded here to + a single value: most callers want the generic ``"molecule"``, but + ``iter_thermo_records`` identifies an unnamed record by its position in + the file (``f"record {position}"``) instead, so that case is still + threaded through explicitly. + """ + return mol.GetProp("_Name") if mol.HasProp("_Name") else default + + def _is_collinear(atoms: ase.Atoms) -> bool: """True if all atoms lie on a single line. @@ -175,7 +190,7 @@ def _symmetry_number(mol: Chem.Mol) -> int: logger.warning( "Molecule %s has an unparseable 'symmetry_number' property " "(%r); falling back to sigma=1.", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), mol.GetProp("symmetry_number"), ) return 1 @@ -198,7 +213,7 @@ def _symmetry_number(mol: Chem.Mol) -> int: "(%d); it must be between 1 and %d (the largest external " "rotational symmetry number of any real molecule, for the " "icosahedral point groups). Falling back to sigma=1.", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), value, _MAX_SYMMETRY_NUMBER, ) @@ -213,7 +228,7 @@ def _symmetry_number(mol: Chem.Mol) -> int: "property (2 for water, 6 for ethane, 12 for benzene) when known. " "(Logged once per run; later molecules defaulting the same way are " "silent.)", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), ) _symmetry_default_warned = True return 1 @@ -301,7 +316,7 @@ def _resolve_multiplicity(mol: Chem.Mol) -> int: logger.warning( "Molecule %s has an unparseable 'multiplicity' property; " "deriving it from the radical-electron count instead.", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), ) else: n_electrons = _electron_count(mol) @@ -311,7 +326,7 @@ def _resolve_multiplicity(mol: Chem.Mol) -> int: "Molecule %s has an invalid 'multiplicity' property (%d); " "multiplicity must be >= 1 (2S+1 for spin S >= 0). " "Deriving it from the radical-electron count instead.", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), value, ) elif value > max_multiplicity: @@ -320,7 +335,7 @@ def _resolve_multiplicity(mol: Chem.Mol) -> int: "a %d-electron species cannot exceed multiplicity %d " "(2S+1 with every electron unpaired). Deriving it from " "the radical-electron count instead.", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), value, n_electrons, max_multiplicity, ) elif value % 2 == n_electrons % 2: @@ -330,7 +345,7 @@ def _resolve_multiplicity(mol: Chem.Mol) -> int: "(2S+1 requires odd multiplicity for an even-electron " "species, even multiplicity for an odd-electron one). " "Deriving it from the radical-electron count instead.", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), value, n_electrons, ) else: @@ -358,7 +373,7 @@ def _resolve_multiplicity(mol: Chem.Mol) -> int: "for an odd-electron one). The drawing may hide an open shell. Set " "the 'multiplicity' property explicitly; the electronic entropy term " "is otherwise wrong by up to R*ln(3) = 0.65 kcal/mol in T*S.", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), multiplicity, n_electrons, ) @@ -380,7 +395,7 @@ def _resolve_multiplicity(mol: Chem.Mol) -> int: "drawing is closed-shell; multiplicity 1 is assumed and the " "electronic entropy term will be wrong. Set the 'multiplicity' " "property explicitly.", - mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule", + _mol_name(mol), ) return multiplicity @@ -1188,7 +1203,7 @@ def do_mol_thermo(mol: Chem.Mol, multiplicity = _resolve_multiplicity(mol) spin = (multiplicity - 1) / 2.0 - name = mol.GetProp("_Name") if mol.HasProp("_Name") else "molecule" + name = _mol_name(mol) # Project translation and rotation out of the Hessian instead of taking # VibrationsData.get_energies()'s raw 3N spectrum and letting # IdealGasThermo guess which entries are vibrations. `atoms` supplies the @@ -1282,7 +1297,7 @@ def do_mol_thermo(mol: Chem.Mol, "mode(s) it was meant to include.", name, n_used, len(vib_e), len(vib_e) - n_used, ) - H = thermo.get_enthalpy(temperature=T) * ev2hatree + H = thermo.get_enthalpy(temperature=T) * EV_TO_HARTREE # ASE's get_entropy returns entropy in eV/K, so this value is Hartree/K, not # Hartree. Name the property accordingly so a downstream G = H - T*S # reconstruction is not off by a factor of T. @@ -1294,14 +1309,14 @@ def do_mol_thermo(mol: Chem.Mol, # (1e5 Pa), so this applies the -kB*T*ln(P/P_ref) correction to report G at # 1 atm -- matching ORCA/Gaussian. The translational-entropy difference vs # 1 bar is R*T*ln(1.01325) = ~0.0078 kcal/mol at 298.15 K. - S = thermo.get_entropy(temperature=T, pressure=STANDARD_PRESSURE) * ev2hatree - G = thermo.get_gibbs_energy(temperature=T, pressure=STANDARD_PRESSURE) * ev2hatree + S = thermo.get_entropy(temperature=T, pressure=STANDARD_PRESSURE) * EV_TO_HARTREE + G = thermo.get_gibbs_energy(temperature=T, pressure=STANDARD_PRESSURE) * EV_TO_HARTREE mol.SetProp("H_hartree", str(H)) mol.SetProp("S_hartree_per_K", str(S)) mol.SetProp("T_K", str(T)) mol.SetProp("G_hartree", str(G)) - mol.SetProp("E_hartree", str(e * ev2hatree)) + mol.SetProp("E_hartree", str(e * EV_TO_HARTREE)) # Only now, with every thermo property computed and set, overwrite mol's # conformer with the relaxed geometry. Deliberately deferred from the top @@ -1498,8 +1513,7 @@ def iter_thermo_records(mols) -> Iterator[Chem.Mol]: logger.warning( "Skipping %s: no 3D conformer, so there is no geometry to " "evaluate.", - mol.GetProp("_Name") if mol.HasProp("_Name") else - f"record {position}", + _mol_name(mol, default=f"record {position}"), ) continue yield mol diff --git a/src/Auto3D/SPE.py b/src/Auto3D/SPE.py index 8d5a2bfa..09ad68ca 100644 --- a/src/Auto3D/SPE.py +++ b/src/Auto3D/SPE.py @@ -8,10 +8,10 @@ from Auto3D.batch_opt.model_wrapper import EnForce_ANI from Auto3D.batch_opt.padding import pad_from_mols +from Auto3D.constants import EV_TO_HARTREE from Auto3D.model_factory import create_model, get_device from Auto3D.models.preflight import resolve_engine_name from Auto3D.torch_config import TorchConfig, configure_torch -from Auto3D.utils.chemistry import hartree2ev from Auto3D.utils.logging_config import get_logger from Auto3D.utils.validation import ( check_engine_supports_molecules, @@ -24,8 +24,6 @@ __all__ = ["calc_spe"] -ev2hatree = 1/hartree2ev - def calc_spe( path: str, @@ -167,7 +165,7 @@ def calc_spe( with Chem.SDWriter(str(outpath)) as f: for i, mol in enumerate(mols): - mol.SetProp('E_hartree', str(es[i] * ev2hatree)) + mol.SetProp('E_hartree', str(es[i] * EV_TO_HARTREE)) f.write(mol) return str(outpath) diff --git a/src/Auto3D/auto3D.py b/src/Auto3D/auto3D.py index 1461d950..c6347d19 100644 --- a/src/Auto3D/auto3D.py +++ b/src/Auto3D/auto3D.py @@ -10,7 +10,6 @@ from pathlib import Path from typing import TYPE_CHECKING -import torch from rdkit import Chem if TYPE_CHECKING: @@ -22,16 +21,14 @@ from Auto3D.config import Auto3DOptions from Auto3D.exceptions import ConfigurationError from Auto3D.isomers import IsomerEngineFactory -from Auto3D.model_factory import create_model +from Auto3D.job_layout import create_chunk_meta_names +from Auto3D.model_factory import create_model, get_device from Auto3D.models.preflight import preflight_model from Auto3D.ranking import ranking -from Auto3D.utils.file_ops import ( - create_chunk_meta_names, - find_smiles_not_in_sdf, - reorder_sdf, - smiles2smi, -) from Auto3D.utils.logging_config import configure_logging, get_logger +from Auto3D.utils.reconciliation import find_smiles_not_in_sdf +from Auto3D.utils.sdf_io import reorder_sdf +from Auto3D.utils.smi_io import smiles2smi from Auto3D.utils.validation import check_input, check_valid_configuration # Pipeline workers live in workflow_workers to break the auto3D<->workflow import @@ -70,7 +67,21 @@ def main( see ``WorkflowOrchestrator._finalize_output``). Raises: - SystemExit: If input validation fails or no structures converge. + ConfigurationError: If path is None, k/window not specified, or the + configuration (including the optimizing engine name) is + otherwise invalid. + FileFormatError: If the input file format is not supported. + OptimizationError: If no structure converged. + ModelLoadError: If the optimizing model could not be obtained or + loaded. + DependencyError: If a required optional dependency is missing. + + None of the above is a ``SystemExit``: ``main()`` itself catches + nothing and lets ``WorkflowOrchestrator.run()``'s exceptions (see + ``WorkflowOrchestrator._validate_input``/``_finalize_output``) + propagate as-is. Only the CLI layer (``cli/errors.handle_error``) + converts them to a process exit code; a direct Python-API caller + gets the exception object. """ # Configure logging based on verbose setting configure_logging(verbose=args.verbose) @@ -208,15 +219,18 @@ def smiles2mols(smiles: list[str], args: Auto3DOptions) -> list[Chem.Mol]: ) isomer_engine.run() - # optimize conformers - if args.use_gpu: - if isinstance(args.gpu_idx, int): - idx = args.gpu_idx - else: - idx = args.gpu_idx[0] - device = torch.device(f"cuda:{idx}") - else: - device = torch.device("cpu") + # optimize conformers. gpu_idx may be a single int or a list (one + # entry per GPU, for main()'s multi-process path); smiles2mols is + # single-process, so only the first index is ever used. Resolved + # through model_factory.get_device -- the single owner of gpu_idx -> + # torch.device -- rather than re-building the `cuda:{idx}` string + # here, which used to bypass get_device's own out-of-range GPUError + # entirely (defense in depth: check_valid_configuration above already + # range-checks gpu_idx for this entry point, but get_device is where + # every other caller, incl. calc_spe/opt_geometry/calc_thermo, gets + # that check). + idx = args.gpu_idx if isinstance(args.gpu_idx, int) else args.gpu_idx[0] + device = get_device(idx, use_gpu=args.use_gpu) opt_config = args.to_optimization_config() # Built in this process, which is also the one that runs the # optimization -- `smiles2mols` is single-process, so there is no spawn diff --git a/src/Auto3D/batch_opt/ANI2xt_no_rep.py b/src/Auto3D/batch_opt/ANI2xt_no_rep.py index e46ec217..b781df8a 100644 --- a/src/Auto3D/batch_opt/ANI2xt_no_rep.py +++ b/src/Auto3D/batch_opt/ANI2xt_no_rep.py @@ -1,10 +1,11 @@ import os +from collections.abc import Sequence import torch import torch.nn as nn from Auto3D.models.species import ANI2XT_INDEX -from Auto3D.utils.chemistry import hartree2ev +from Auto3D.utils.energy import hartree2ev # Note: Do NOT set torch.manual_seed() at module level. # Random seed should be controlled by the caller, not by importing a module. @@ -70,6 +71,134 @@ def _atomic_mlp(aev_dim: int, widths: tuple[int, int, int]) -> nn.Sequential: ) +#: Number of per-element networks, i.e. ``len(WIDTHS)``. Exposed so callers can +#: size the per-element index list without reaching into a (possibly +#: ``torch.compile``-wrapped) module. +NUM_ELEMENTS: int = len(WIDTHS) + + +def element_indices(species_idx: torch.Tensor, num_elements: int = NUM_ELEMENTS) -> list[torch.Tensor]: + """Row indices of each element's atoms on the FLATTENED ``(batch*atoms,)`` axis. + + Returns exactly what ``[torch.nonzero(flat == e)[0] for e in range(n)]`` + returns -- same indices, same ascending order -- but with **one** host + readback instead of ``n`` of them. On CUDA each ``nonzero`` is a + synchronization (its output shape is data-dependent, so ATen must copy the + match count to the host), and this loop ran on every ANI2xt forward. + + The construction: bucket each atom (``0`` for padded/negative species, + ``1..n`` for elements ``0..n-1``, ``n+1`` for anything out of range), stable + ``argsort`` into element order, count buckets with a fixed-size + ``scatter_add_``, then ``split`` the sorted order by those counts. Only the + counts cross to the host. Out-of-range species get their own bucket and are + discarded rather than clamped into a neighbouring element's network, which is + what ``flat == e`` did. + + Ordering is *not* load-bearing for energies -- each network is applied row by + row and ``index_copy`` targets unique rows -- but matching ``nonzero`` + exactly makes bit-identity with the previous implementation checkable, which + ``tests/test_ani2xt_atom_energies.py`` asserts over 200+ species patterns. + + Args: + species_idx: 0-based element indices, ``(batch, atoms)``. Padded slots + carry ``-1``; any negative value is treated as padding. + num_elements: Number of per-element networks. + + Returns: + ``num_elements`` int64 index tensors, ascending, into the flattened axis. + """ + flat = species_idx.reshape(-1) + # Bucket 0 absorbs padding (species < 0); bucket num_elements+1 absorbs + # out-of-range species. Both are dropped from the returned list. + bucket = (flat + 1).clamp(0, num_elements + 1) + order = torch.argsort(bucket, stable=True) + counts = torch.zeros(num_elements + 2, dtype=torch.long, device=flat.device) + counts.scatter_add_(0, bucket, torch.ones_like(bucket)) + sizes = counts.tolist() # the ONE host readback + return list(order.split(sizes))[1:num_elements + 1] + + +def self_atomic_energies( + species_idx: torch.Tensor, + energy_shifts: torch.Tensor, + num_elements: int = NUM_ELEMENTS, +) -> torch.Tensor: + """Per-molecule self-atomic energy shifts, in Hartree. + + A pure function of ``species_idx``, so it is constant for a whole bucket of + conformers and does not belong in the hot path -- it used to be recomputed on + every forward, costing roughly ``4 * num_elements`` kernel launches per step + for a value that never changed. Precompute it once and pass it to + ``ANI2xt.forward``. + + Summation order over elements is preserved, so the result is bit-identical to + the inline version it replaced. + + Args: + species_idx: 0-based element indices, ``(batch, atoms)``. + energy_shifts: Per-element shift, ``(num_elements,)``, float64. + num_elements: Number of per-element networks. + + Returns: + float64 tensor, ``(batch,)``. + """ + out = torch.zeros(species_idx.shape[0], device=species_idx.device, dtype=torch.float64) + for elem_idx in range(num_elements): + counts = (species_idx == elem_idx).sum(dim=1).to(torch.float64) + out += counts * energy_shifts[elem_idx] + return out + + +def _atom_energies( + networks: nn.ModuleList | Sequence[nn.Module], + aev_flat: torch.Tensor, + elem_index: list[torch.Tensor], + n_rows: int, +) -> torch.Tensor: + """Per-atom energies over a flattened atom axis, with no data-dependent shapes. + + Module-level and taking ``networks`` as a parameter so it is testable without + torchani (the AEV computer is the only part of ANI2xt that needs it, and it + runs before this). + + Three things happen here, all of which matter: + + * **The ``if mask.any():`` guard is gone.** It protected nothing: + ``network(empty)`` returns an empty tensor and ``index_copy`` with an empty + index is a no-op, which is why deleting it is bit-identical even for a + batch containing only 2 of the 7 elements. It cost 7 host-device + synchronizations per forward, and it made this whole frame uncompilable -- + a data-dependent branch *inside* a ``for`` loop gives Dynamo nowhere to + place a resume point, so it skipped the frame entirely and + ``torch.compile`` produced **zero** subgraphs for ``ANI2xt.forward``. + * **Indices are passed in, not computed here.** ``nonzero`` and boolean-mask + indexing are dynamic-output-shape ops, so computing them in this loop would + graph-break too and the frame would still be skipped. Only a loop body with + no data-dependent op at all compiles: this form is **one** subgraph and + passes ``fullgraph=True``. + * **Functional ``index_copy``, not ``index_copy_``.** The out-of-place form + avoids input-mutation handling under ``torch.compile``, and is what was + measured at one subgraph. + + Args: + networks: One energy network per element, in ``WIDTHS`` order. + aev_flat: AEV features, ``(n_rows, aev_dim)``. + elem_index: Per-element int64 row indices, from :func:`element_indices`. + n_rows: ``batch * atoms``; passed explicitly so the output shape never + depends on a tensor value. + + Returns: + float64 per-atom energies, ``(n_rows,)``. Rows belonging to no element + (padded or out-of-range) stay zero. + """ + out = torch.zeros(n_rows, dtype=torch.float64, device=aev_flat.device) + for elem_idx, network in enumerate(networks): + idx = elem_index[elem_idx] + selected = aev_flat.index_select(0, idx) + out = out.index_copy(0, idx, network(selected).squeeze(-1).to(torch.float64)) + return out + + class ANI2xt(nn.Module): def __init__(self, device, state_dict=ani_2xt_dict, periodic_table_index=False): super().__init__() @@ -121,7 +250,7 @@ def __init__(self, device, state_dict=ani_2xt_dict, periodic_table_index=False): # Canonical atomic-number -> ANI2xt species index map (species.ANI2XT_INDEX). self.periodict2idx = dict(ANI2XT_INDEX) - def forward(self, species, coords): + def forward(self, species, coords, elem_index=None, self_energies=None): """Compute molecular energies. Args: @@ -129,10 +258,31 @@ def forward(self, species, coords): If periodic_table_index=True, uses atomic numbers (1=H, 6=C, etc.) Otherwise, uses sequential indices (0=H, 1=C, 2=N, 3=O, 4=F, 5=S, 6=Cl) coords: Tensor of shape (batch, num_atoms, 3) with atomic coordinates + elem_index: Optional per-element row indices on the flattened atom + axis, as returned by :func:`element_indices`. Species are + constant for a bucket of conformers while coordinates are not, + so a caller that optimizes the same molecules for many steps + should compute this once and pass it in. ``None`` computes it + here, which keeps every existing caller working unchanged but + pays one host readback per forward -- and, because the + computation has a data-dependent output shape, reintroduces the + graph break that stops ``torch.compile`` from compiling this + method at all. + self_energies: Optional per-molecule self-atomic energy shifts in + Hartree, as returned by :func:`self_atomic_energies`. Also a + pure function of ``species``. ``None`` computes it here. Returns: Tensor of shape (batch,) with molecular energies in eV + Raises: + ValueError: ``elem_index`` or ``self_energies`` was supplied while + ``periodic_table_index=True``. The caller would have computed + them from atomic numbers rather than from the remapped 0-based + indices this model's networks are indexed by, and the mistake is + silent -- atomic number 6 (carbon) is a valid *index* for + chlorine. + Note: The AEV/network path runs in float32 (coords dtype), but the per-atom and self-atomic energies are accumulated in float64 so the @@ -145,6 +295,14 @@ def forward(self, species, coords): float32 coords), matching the AIMNet2 adapter's output contract. """ if self.periodic: + if elem_index is not None or self_energies is not None: + raise ValueError( + "ANI2xt.forward: elem_index/self_energies are indexed by " + "0-based network index, but this instance was built with " + "periodic_table_index=True and receives atomic numbers. " + "Precompute them from the remapped species, or let forward " + "compute them." + ) # Convert atomic numbers to sequential indices species_idx = species.clone() for key, val in self.periodict2idx.items(): @@ -156,36 +314,33 @@ def forward(self, species, coords): # padder). -1 is not in periodict2idx, so it survives unchanged here and # is passed to the AEV computer, which relies on TorchANI's convention # that a species index of -1 marks a dummy/masked atom (excluded from the - # AEV and the per-element energy loop below, where no elem_idx == -1). + # AEV and from every per-element index below, which drop negatives). # This correctness depends on -1 being TorchANI's masked-atom sentinel. # Compute AEVs (use new API: aev_computer(species, coords)) aev = self.aev_computer(species_idx, coords) # (batch, num_atoms, aev_dim) - # Compute per-atom energies. Accumulate in float64 so the float64 - # energy_shifts buffer is meaningful; the network output is float32 and - # is cast up explicitly (an fp64-dest, fp32-source index_put would raise). batch_size, num_atoms = species_idx.shape - atom_energies = torch.zeros(batch_size, num_atoms, device=coords.device, dtype=torch.float64) - - for elem_idx, network in enumerate(self.networks): - # Find atoms of this element type - mask = (species_idx == elem_idx) - if mask.any(): - # Get AEVs for atoms of this element - elem_aev = aev[mask] # (num_elem_atoms, aev_dim) - # Compute atomic energies - elem_energies = network(elem_aev).squeeze(-1) # (num_elem_atoms,) - atom_energies[mask] = elem_energies.to(torch.float64) + n_rows = batch_size * num_atoms + if elem_index is None: + elem_index = element_indices(species_idx, len(self.networks)) + if self_energies is None: + self_energies = self_atomic_energies( + species_idx, self.energy_shifts, len(self.networks)) + + # Per-atom energies, accumulated in float64 so the float64 energy_shifts + # buffer is meaningful; the network output is float32 and is cast up + # explicitly (index_copy will not cast for us, in either direction). + # reshape with an explicit trailing size rather than -1: inferring -1 on + # a zero-element tensor is ambiguous and raises. + atom_energies = _atom_energies( + self.networks, + aev.reshape(n_rows, aev.shape[-1]), + elem_index, + n_rows, + ) # Sum per-atom energies to get molecular energies - atomic_energies = atom_energies.sum(dim=1) # (batch,) - - # Add self-energies (energy shifts) - self_energies = torch.zeros(batch_size, device=coords.device, dtype=torch.float64) - for elem_idx in range(len(self.networks)): - mask = (species_idx == elem_idx) - counts = mask.sum(dim=1).to(torch.float64) # (batch,) - self_energies += counts * self.energy_shifts[elem_idx] + atomic_energies = atom_energies.reshape(batch_size, num_atoms).sum(dim=1) # (batch,) # Total energy in Hartree, convert to eV total_energy = (atomic_energies + self_energies) * hartree2ev diff --git a/src/Auto3D/batch_opt/batchopt.py b/src/Auto3D/batch_opt/batchopt.py index 8d4ff6e4..4a49a587 100644 --- a/src/Auto3D/batch_opt/batchopt.py +++ b/src/Auto3D/batch_opt/batchopt.py @@ -386,7 +386,7 @@ def run(self): # object, so no atom mapping is needed. This covers the neural # network optimization step only; clash relief (a separate, # earlier force-field relaxation) is guarded at its own call - # site in Auto3D.utils.chemistry.relieve_clash. + # site in Auto3D.clash_relief.relieve_clash. if not apply_optimized_coords(mol, coords_out[i]): n_stereo_changed += 1 f.write(mol) @@ -397,7 +397,7 @@ def run(self): # is attached to (see Auto3D.workflow_workers), so a warning # through the module logger never reaches the run log. Emit # through logging.getLogger("auto3d") directly instead -- the - # same fix Auto3D.utils.chemistry.relieve_clash already uses -- + # same fix Auto3D.clash_relief.relieve_clash already uses -- # because this count is documented as user-visible in the log # (CHANGELOG.md, docs/source/migration-4.0.rst). logging.getLogger("auto3d").warning( diff --git a/src/Auto3D/batch_opt/fire_optimizer.py b/src/Auto3D/batch_opt/fire_optimizer.py index 209084d1..3335dbc5 100644 --- a/src/Auto3D/batch_opt/fire_optimizer.py +++ b/src/Auto3D/batch_opt/fire_optimizer.py @@ -162,21 +162,41 @@ def __call__(self, coord: torch.Tensor, forces: torch.Tensor) -> torch.Tensor: return coord + dr - def clean(self, mask: torch.Tensor) -> bool: - """Subset optimizer state to keep only specified molecules. + def clean(self, idx: torch.Tensor) -> bool: + """Subset optimizer state to keep only the molecules at ``idx``. - This method is used to remove converged molecules from the optimization - batch, reducing computation in subsequent steps. + Used to drop molecules that have left the active set, so subsequent + steps do less work. + + Takes an **int64 index**, not a boolean mask. Four boolean-mask reads + here were four host-device synchronizations per optimization step on + CUDA (ATen expands each mask via ``nonzero()`` and copies the element + count to the host to size the output). ``index_select`` with an index + the caller already computed is sync-free, so the caller pays one + ``nonzero`` and this method pays none. Over 2000 steps per bucket that + is 8000 serialization points removed from a loop that does no host work. Args: - mask: Boolean tensor of shape (batch,). True values indicate - molecules to keep in the optimization. + idx: int64 tensor of row indices to keep, shape ``(kept,)``. Must be + int64: a boolean mask would *not* fail loudly here -- ``index_select`` + would reinterpret ``tensor([True, False])`` as indices ``[1, 0]`` + and silently permute the optimizer state -- so the dtype is + checked instead. Returns: Always returns True to indicate success. + + Raises: + ValueError: ``idx`` is not int64. """ - self.v = self.v[mask] - self.Nsteps = self.Nsteps[mask] - self.dt = self.dt[mask] - self.a = self.a[mask] + if idx.dtype != torch.long: + raise ValueError( + "FIRE.clean expects an int64 tensor of row indices to keep, not " + "a boolean mask. Pass torch.nonzero(mask, as_tuple=True)[0]. " + "A bool tensor would be silently reinterpreted as indices." + ) + self.v = self.v.index_select(0, idx) + self.Nsteps = self.Nsteps.index_select(0, idx) + self.dt = self.dt.index_select(0, idx) + self.a = self.a.index_select(0, idx) return True diff --git a/src/Auto3D/batch_opt/optimization_engine.py b/src/Auto3D/batch_opt/optimization_engine.py index 71dab5a9..289a8f2c 100644 --- a/src/Auto3D/batch_opt/optimization_engine.py +++ b/src/Auto3D/batch_opt/optimization_engine.py @@ -3,11 +3,42 @@ This module contains the main optimization loop (n_steps) and status reporting (print_stats) functions extracted from batchopt.py for better modularity. + +Host-device synchronization +--------------------------- +``n_steps`` runs up to 2000 iterations per bucket, so a per-step +host-device serialization point costs 2000 of them. Every subset read and write +in the loop therefore goes through **one** ``torch.nonzero`` per step, whose +int64 result feeds ``index_select`` (reads) and ``index_copy_`` (writes) -- +neither of which synchronizes. + +The loop used to subset with boolean masks instead, which on CUDA is 18 syncs +per step: ATen has to ``nonzero()`` each mask and copy the element count to the +host to size the output, once per masked read and once per masked write, six of +each here plus four more inside ``FIRE.clean``. Reusing a single ``nonzero`` +result across all twelve gathers and scatters takes that to **2 per step** -- +one for the active subset, one inside the step for ``FIRE.clean``, whose mask is +indexed within the active subset rather than the full batch. + +Two, not zero: ``nonzero`` is *itself* the sync, which is precisely why boolean +masking synced. The win is amortizing one over twelve uses, not eliminating it. +``tests/test_optimization_engine_indexing.py`` counts this on CPU (the sync is a +property of the operator, not the device) and fails if boolean-mask indexing +reappears. + +``index_copy_`` is stricter than masked assignment about dtype: mismatched +writes raise rather than cast. Every write below therefore casts explicitly to +the destination dtype. That is not merely tidiness -- it is what lets a custom +NNP return float64 forces, which the boolean-mask loop rejected with +``"Index put requires the source and destination dtypes match"`` whenever two or +more molecules reduced their force in the same step (with exactly one, the value +had ``numel() == 1``, hit ATen's ``masked_fill_`` fast path and silently cast, +so the crash was batch-size dependent). """ from __future__ import annotations from collections.abc import Callable -from typing import Any +from typing import Any, NamedTuple import numpy as np import torch @@ -83,6 +114,247 @@ def print_stats(state: dict[str, Any], patience: int) -> None: (num_total, num_converged, num_dropped, num_active)) +def _emit_progress( + state: dict[str, Any], + patience: int, + progress_cb: Callable[[dict], None] | None, + istep: int, +) -> None: + """Emit one live-progress event, if a callback was supplied. + + Only active when the caller passed a callback (i.e. interactive + ``auto3d run``), so the default library path adds nothing -- including the + two host-device syncs ``optimization_counts`` performs. Wrapped so a + progress-display hiccup can never abort an optimization that is otherwise + fine. + + Args: + state: Optimization state dictionary. + patience: Oscillation patience, needed to split converged from dropped. + progress_cb: Callback, or None to do nothing. + istep: Step number to report. + """ + if progress_cb is None: + return + try: + total, converged, dropped, active = optimization_counts(state, patience) + progress_cb({"step": istep, "total": total, "converged": converged, + "dropped": dropped, "active": active}) + except Exception: + pass + + +class _StepResult(NamedTuple): + """Per-molecule results for one step, ordered by position in the active subset. + + Every field has leading dimension ``active_idx.numel()``, so + ``_scatter_back`` can write all six with the same index. + + Attributes: + coord: Post-step coordinates, ``(active, n_atoms, 3)``. + energy: Energy at the *pre*-step geometry, ``(active,)``. + fmax: Maximum force at the pre-step geometry, ``(active,)``. + still_active: True for molecules that remain in the active set, i.e. + neither force-converged nor dropped as oscillating, ``(active,)``. + smallest_fmax: Running per-molecule force minimum, ``(active, 1)``. + oscillating_count: Steps since that minimum last improved, ``(active,)``. + """ + + coord: torch.Tensor + energy: torch.Tensor + fmax: torch.Tensor + still_active: torch.Tensor + smallest_fmax: torch.Tensor + oscillating_count: torch.Tensor + + +def _step_active_subset( + state: dict[str, Any], + optimizer: FIRE, + active_idx: torch.Tensor, + smallest_fmax0: torch.Tensor, + opttol: float, + patience: int, +) -> _StepResult: + """Take one FIRE step over the still-active molecules. + + Gathers the active rows with ``index_select`` (no sync), evaluates the model, + applies the force-convergence and oscillation criteria, steps the optimizer + and subsets the optimizer's own state. Mutates ``optimizer`` in place; + everything else is returned for ``_scatter_back`` to write. + + Args: + state: Optimization state dictionary. Read, not written. + optimizer: The FIRE optimizer, whose per-molecule state is subset here. + active_idx: int64 row indices of still-active molecules, ascending. + smallest_fmax0: Full-batch running force minimum, ``(batch, 1)``. + opttol: Force convergence tolerance in eV/Angstrom. + patience: Steps without force decrease before dropping as oscillating. + + Returns: + A ``_StepResult`` aligned with ``active_idx``. + """ + coord = state['coord'].index_select(0, active_idx) + numbers = state['numbers'].index_select(0, active_idx) + charges = state['charges'].index_select(0, active_idx) + atom_mask_subset = state['atom_mask'].index_select(0, active_idx) + smallest_fmax = smallest_fmax0.index_select(0, active_idx) + oscillating_count = state['oscillating_count'].index_select(0, active_idx) + + coord.requires_grad_(True) + # atom_mask goes to the model too, not just to the force reduction below: + # an adapter that has to flatten a padded batch (AIMNet2) needs to know + # which slots are real, and deriving that from a species sentinel is what + # audit C13 forbids. + e, f = state['nn'].forward_batched( + coord, numbers, charges, + atom_mask=atom_mask_subset, + ) # Key step to calculate all energies and forces. + coord.requires_grad_(False) + + # Zero forces on padded atom slots so convergence is independent of how the + # model treats ghost atoms. atom_mask is True for real atoms. Deriving this + # from a sentinel value (numbers == species_pad) broke for any model whose + # species_pad collides with a real index (audit C13). Masking before the + # optimizer step also keeps padded atoms from drifting. + pad_mask = ~atom_mask_subset.unsqueeze(-1) + f = f.masked_fill(pad_mask, 0.0) + # Norm is the length of each force vector; the max over atoms is the + # per-molecule convergence measure. + fmax = f.norm(dim=-1).max(dim=-1)[0] + + # The force-convergence test runs BEFORE the FIRE step, and must stay there. + not_converged_post1 = fmax > opttol + # Detach the optimizer output so the next step starts from a leaf tensor. + # Production adapters return detached forces, so coord never tracks grad + # across steps; detaching here makes the loop robust to NNPs whose forces + # still carry a grad graph (otherwise the next requires_grad_ would error). + stepped = optimizer(coord, f).detach() + # A structure that has just met the force criterion keeps the geometry its + # force was measured at; only the ones still moving take the step. + # + # The convergence test above and the step used to run in the other order, + # so every structure took one more FIRE step *after* the force that + # declared it converged. `Converged` then described the geometry before + # that step while the reported `fmax` (recomputed at the end of n_steps) + # described the geometry after it -- and `batchopt` writes both onto the + # same record. A consumer filtering on `fmax <= opt_tol` and one filtering + # on `Converged == "True"` got different sets from one file. Measured on a + # hermetic harmonic potential before this change: fmax up to 6.9x the + # tolerance beside `Converged=True`. The discrepancy grows with stiffness, + # so a soft test case hides it entirely -- which is how it survived being + # looked at twice. + # + # Structures leaving the active set as oscillating are stepped like any + # other: they are reported `Converged=False`, so no consistency is claimed + # for them, and n_steps' end-of-function recompute makes their `fmax` match + # their coordinates regardless. + coord = torch.where(not_converged_post1.view(-1, 1, 1), stepped, coord) + + # Update the running force minimum and the oscillation counter. Both are + # per-molecule elementwise updates that need no index at all, so they are + # `torch.where` rather than masked assignment. + # + # `torch.where`, NOT `torch.minimum`: `<` is False for NaN, so the masked + # assignment this replaced *kept* the previous smallest_fmax when a + # molecule's force went NaN, whereas `minimum` propagates the NaN and would + # poison that molecule's oscillation tracking for the rest of the run. + # `_validate_outputs` should make NaN unreachable through the adapters, but + # "should be unreachable" is not a reason to change semantics. + fmax_col = fmax.reshape(-1, 1) + fmax_reduced = (fmax_col < smallest_fmax).reshape(-1, ) + smallest_fmax = torch.where(fmax_reduced.unsqueeze(-1), fmax_col, smallest_fmax) + # Reduced -> reset to 0; not reduced -> increment. This is one `where` in + # place of "zero the reduced entries, then add ~reduced", whose ordering was + # load-bearing: a reduced molecule was zeroed and then incremented by False. + oscillating_count = torch.where( + fmax_reduced, torch.zeros_like(oscillating_count), oscillating_count + 1) + not_oscillating = oscillating_count < patience + + # Combine the convergence criteria. An `& ~energy_converged` term stood + # here until 4.0.0; `energy_converged` required `fmax < opttol` while + # `not_converged_post1` is `fmax > opttol`, so the term was the identity + # of `&` wherever it was consulted and false-dominated elsewhere -- it + # could never change an outcome, including at the `fmax == opttol` + # boundary where both comparisons are false (audit M1). + still_active = not_converged_post1 & not_oscillating + + # Second nonzero of the step: `still_active` is indexed within the active + # subset, not the full batch, so it cannot reuse `active_idx`. FIRE.clean + # takes an int64 index precisely so this is one sync instead of four. + optimizer.clean(torch.nonzero(still_active, as_tuple=True)[0]) + + return _StepResult( + coord=coord, + energy=e.detach(), + fmax=fmax, + still_active=still_active, + smallest_fmax=smallest_fmax, + oscillating_count=oscillating_count, + ) + + +def _scatter_back( + state: dict[str, Any], + active_idx: torch.Tensor, + smallest_fmax0: torch.Tensor, + result: _StepResult, +) -> None: + """Write one step's results back into the full-batch state tensors. + + All six writes reuse ``active_idx``, so they cost no synchronization. + ``index_copy_`` requires the source dtype to match the destination exactly + (it will not cast, in either direction), so each source is cast at the call + site -- which is also the only reason a custom NNP returning float64 forces + works here. + + Args: + state: Optimization state dictionary, mutated in place. + active_idx: int64 row indices the results correspond to. + smallest_fmax0: Full-batch running force minimum, mutated in place. + result: The values returned by ``_step_active_subset``. + """ + # Converged structures are excluded from every subsequent step. + state['converged_mask'].index_copy_( + 0, active_idx, (~result.still_active).to(state['converged_mask'].dtype)) + state['fmax'].index_copy_( + 0, active_idx, result.fmax.to(state['fmax'].dtype)) + state['energy'].index_copy_( + 0, active_idx, result.energy.to(state['energy'].dtype)) + state['coord'].index_copy_( + 0, active_idx, result.coord.to(state['coord'].dtype)) + smallest_fmax0.index_copy_( + 0, active_idx, result.smallest_fmax.to(smallest_fmax0.dtype)) + state['oscillating_count'].index_copy_( + 0, active_idx, result.oscillating_count.to(state['oscillating_count'].dtype)) + + +def _recompute_final_energy_and_fmax(state: dict[str, Any]) -> None: + """Re-evaluate energy and fmax at the final reported geometry. + + Energy and fmax stored during the loop are evaluated at the *pre*-step + geometry, while the stored coordinates are post-step (the loop always takes + one FIRE step after measuring forces). Recompute both once at the end so + ``state['energy']`` and ``state['fmax']`` correspond to the reported + ``state['coord']``. The adapters differentiate internally for forces, so + grad must be enabled; the forces themselves were previously discarded. + + Args: + state: Optimization state dictionary, mutated in place. + """ + final_coord = state['coord'].detach().clone().requires_grad_(True) + e_final, f_final = state['nn'].forward_batched( + final_coord, state['numbers'], state['charges'], + atom_mask=state['atom_mask'], + ) + state['energy'] = e_final.detach().to(state['energy'].dtype) + # Zero padded-atom force slots before the reduction, matching the in-loop + # convergence check, so reported fmax is independent of how the model treats + # ghost atoms. atom_mask is True for real atoms (audit C13). + f_final = f_final.detach().masked_fill(~state['atom_mask'].unsqueeze(-1), 0.0) + state['fmax'] = f_final.norm(dim=-1).max(dim=-1)[0].to(state['fmax'].dtype) + + def n_steps( state: dict[str, Any], n: int, @@ -106,6 +378,11 @@ def n_steps( the structure, so the term was the identity of ``&`` at every element (see ``test_convergence_outcome_never_depends_on_energy_stability``). + The per-step work is split across ``_step_active_subset`` (gather, model, + criteria, FIRE step) and ``_scatter_back`` (the writes), both driven by a + single ``torch.nonzero`` of the active set -- see the module docstring for + why that matters and what it costs. + Args: state: Optimization state dictionary containing: - numbers: Atomic numbers, shape (batch, n_atoms) @@ -130,9 +407,10 @@ def n_steps( convention Auto3D itself uses for ANI2xt, where 0 is hydrogen (audit C13). Defaults to None, which is a no-op (every atom treated as real) for unpadded batches or any caller that omits it. + progress_cb: Optional callback invoked every 10 steps and once at the + end with a dict of step/total/converged/dropped/active counts. """ numbers = state['numbers'] - charges = state['charges'] coord = state['coord'] if atom_mask is None: @@ -147,13 +425,11 @@ def n_steps( # The following two terms are used to detect oscillating conformers smallest_fmax0 = torch.tensor(np.ones((len(coord), 1)) * 999, dtype=torch.float).to(coord.device) - # Integer step counter: only ever incremented by a bool mask and compared - # to `patience`, so use torch.long rather than float for a quantity that is - # conceptually an integer count. - oscillating_count0 = torch.zeros(len(coord), dtype=torch.long, - device=coord.device) - - state["oscillating_count"] = oscillating_count0 + # Integer step counter: only ever compared to `patience` and reset to zero, + # so use torch.long rather than float for a quantity that is conceptually + # an integer count. + state["oscillating_count"] = torch.zeros(len(coord), dtype=torch.long, + device=coord.device) istep = 0 # Initialize in case loop doesn't execute (n=0) # Plain range, not tqdm. A bar over `range(1, n+1)` measures the *step @@ -174,143 +450,37 @@ def n_steps( if istep % 10 == 0 and not not_converged.any(): break - coord = state['coord'][not_converged] # Subset coordinates, size=not_converged. + # THE one sync of the gather half of the step. Its int64 result is + # reused by six index_select reads and six index_copy_ writes; the + # boolean-mask spelling would have paid a nonzero() for each of the + # twelve. Indices come out ascending, so gathered row order -- and + # therefore every downstream reduction -- is unchanged. + active_idx = torch.nonzero(not_converged, as_tuple=True)[0] # On non-throttle steps we may reach here after every molecule has - # converged (the .any() break only runs every 10 steps). Subsetting then - # yields a zero-length batch; bail out before the (empty) NN call rather - # than feeding an empty batch through the model. `.shape[0]` is host-side - # tensor metadata, so this guard adds no host-device sync. - if coord.shape[0] == 0: + # converged (the .any() break only runs every 10 steps). Stepping then + # would feed a zero-length batch through the model, so bail out first. + # `numel()` on the nonzero result is host-side metadata that the nonzero + # already made available, so this guard adds no further sync. + if active_idx.numel() == 0: break - numbers = state['numbers'][not_converged] - charges = state['charges'][not_converged] - atom_mask_subset = state['atom_mask'][not_converged] - smallest_fmax = smallest_fmax0[not_converged] - oscillating_count = state["oscillating_count"][not_converged] - - coord.requires_grad_(True) - # atom_mask goes to the model too, not just to the force reduction - # below: an adapter that has to flatten a padded batch (AIMNet2) needs - # to know which slots are real, and deriving that from a species - # sentinel is what audit C13 forbids. - e, f = state['nn'].forward_batched( - coord, numbers, charges, - atom_mask=atom_mask_subset, - ) # Key step to calculate all energies and forces. - coord.requires_grad_(False) - - # Zero forces on padded atom slots so convergence is independent of how - # the model treats ghost atoms. atom_mask is True for real atoms. - # Deriving this from a sentinel value (numbers == species_pad) broke - # for any model whose species_pad collides with a real index (audit - # C13). Use the loop-local `atom_mask_subset` (state['atom_mask'][not_converged]) - # so the mask aligns with the current batch of f. Masking before the - # optimizer step also keeps padded atoms from drifting. - pad_mask = ~atom_mask_subset.unsqueeze(-1) - f = f.masked_fill(pad_mask, 0.0) - fmax = f.norm(dim=-1).max(dim=-1)[ - 0] # Tensor, Norm is the length of each vector. Here it returns the maximum force length for each conformer. Size (100) - not_converged_post1 = fmax > opttol - # Detach the optimizer output so the next step starts from a leaf tensor. - # Production adapters return detached forces, so coord never tracks grad - # across steps; detaching here makes the loop robust to NNPs whose forces - # still carry a grad graph (otherwise the next requires_grad_ would error). - stepped = optimizer(coord, f).detach() - # A structure that has just met the force criterion keeps the geometry its - # force was measured at; only the ones still moving take the step. - # - # The convergence test above and the step used to run in the other order, - # so every structure took one more FIRE step *after* the force that - # declared it converged. `Converged` then described the geometry before - # that step while the reported `fmax` (recomputed at the end of this - # function) described the geometry after it -- and `batchopt` writes both - # onto the same record. A consumer filtering on `fmax <= opt_tol` and one - # filtering on `Converged == "True"` got different sets from one file. - # Measured on a hermetic harmonic potential before this change: fmax up to - # 6.9x the tolerance beside `Converged=True`. The discrepancy grows with - # stiffness, so a soft test case hides it entirely -- which is how it - # survived being looked at twice. - # - # Structures leaving the active set as oscillating are stepped like any - # other: they are reported `Converged=False`, so no consistency is claimed - # for them, and the end-of-function recompute makes their `fmax` match - # their coordinates regardless. - coord = torch.where(not_converged_post1.view(-1, 1, 1), stepped, coord) - - # Update smallest_fmax for each molecule - fmax_reduced = fmax.reshape(-1, 1) < smallest_fmax - fmax_reduced = fmax_reduced.reshape(-1, ) - smallest_fmax[fmax_reduced] = fmax.reshape(-1, 1)[fmax_reduced] - # Reduce count to 0 for reducing; raise count for non-reducing - oscillating_count[fmax_reduced] = 0 - fmax_not_reduced = ~fmax_reduced - oscillating_count += fmax_not_reduced - not_oscillating = oscillating_count < patience - - # Combine the convergence criteria. An `& ~energy_converged` term stood - # here until 4.0.0; `energy_converged` required `fmax < opttol` while - # `not_converged_post1` is `fmax > opttol`, so the term was the identity - # of `&` wherever it was consulted and false-dominated elsewhere -- it - # could never change an outcome, including at the `fmax == opttol` - # boundary where both comparisons are false (audit M1). - not_converged_post = not_converged_post1 & not_oscillating - - optimizer.clean(not_converged_post) # Subset v, a in FIRE for next optimization - - state['converged_mask'][ - not_converged] = ~ not_converged_post # Update converged_mask, so that converged structures will not be updated in future steps. - state['fmax'][ - not_converged] = fmax # Update fmax for conformers that are optimized in this iteration - state['energy'][ - not_converged] = e.detach().to(state['energy'].dtype) # Update energy for conformers that are optimized in this iteration - state['coord'][ - not_converged] = coord # Update coordinates for conformers that are optimized in this iteration - smallest_fmax0[not_converged] = smallest_fmax # Update smallest_fmax for each conformer - state["oscillating_count"][ - not_converged] = oscillating_count # Update counts for continuous no reduction in fmax + + result = _step_active_subset( + state, optimizer, active_idx, smallest_fmax0, opttol, patience) + _scatter_back(state, active_idx, smallest_fmax0, result) # Print stats every 10% of steps (avoid division by zero for small n) if n >= 10 and (istep % (n // 10)) == 0: print_stats(state, patience) - # Emit a live-progress event for the CLI display. Reuses the istep % 10 - # sync cadence; only active when a callback was supplied (i.e. interactive - # `auto3d run`), so the default/library path adds nothing. Guarded so a - # progress hiccup can never abort the optimization. - if progress_cb is not None and istep % 10 == 0: - try: - total, converged, dropped, active = optimization_counts(state, patience) - progress_cb({"step": istep, "total": total, "converged": converged, - "dropped": dropped, "active": active}) - except Exception: - pass + # Live-progress event for the CLI display, on the same 10-step cadence + # as the throttled sync above. + if istep % 10 == 0: + _emit_progress(state, patience, progress_cb, istep) # Final event so the display reflects the converged end state. - if progress_cb is not None: - try: - total, converged, dropped, active = optimization_counts(state, patience) - progress_cb({"step": istep, "total": total, "converged": converged, - "dropped": dropped, "active": active}) - except Exception: - pass - - # Energy and fmax stored during the loop are evaluated at the pre-step - # geometry, while the stored coordinates are post-step (the loop always takes - # one FIRE step after measuring forces). Recompute both once at the final - # geometry so state['energy'] and state['fmax'] correspond to the reported - # state['coord']. The adapters differentiate internally for forces, so grad - # must be enabled; the forces were previously discarded. - final_coord = state['coord'].detach().clone().requires_grad_(True) - e_final, f_final = state['nn'].forward_batched( - final_coord, state['numbers'], state['charges'], - atom_mask=state['atom_mask'], - ) - state['energy'] = e_final.detach().to(state['energy'].dtype) - # Zero padded-atom force slots before the reduction, matching the in-loop - # convergence check, so reported fmax is independent of how the model treats - # ghost atoms. atom_mask is True for real atoms (audit C13). - f_final = f_final.detach().masked_fill(~state['atom_mask'].unsqueeze(-1), 0.0) - state['fmax'] = f_final.norm(dim=-1).max(dim=-1)[0].to(state['fmax'].dtype) + _emit_progress(state, patience, progress_cb, istep) + + _recompute_final_energy_and_fmax(state) if istep == (n): logger.info("Reaching maximum optimization step:") diff --git a/src/Auto3D/chunk_manager.py b/src/Auto3D/chunk_manager.py index d533436d..6d3bd7b8 100644 --- a/src/Auto3D/chunk_manager.py +++ b/src/Auto3D/chunk_manager.py @@ -13,8 +13,8 @@ import psutil import torch -from Auto3D.utils.file_ops import SDF2chunks from Auto3D.utils.logging_config import get_logger +from Auto3D.utils.sdf_io import SDF2chunks if TYPE_CHECKING: import logging diff --git a/src/Auto3D/clash_relief.py b/src/Auto3D/clash_relief.py new file mode 100644 index 00000000..b2ec2cd1 --- /dev/null +++ b/src/Auto3D/clash_relief.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python +"""Force-field relief of a clashing embedded conformer. + +A domain operation, not a generic helper: it runs MMFF (with a UFF fallback), +it can change the molecule it is handed, and it makes a keep/reject decision +about stereochemistry. That is why it sits beside the embedding code rather +than under ``Auto3D.utils``, whose modules are leaves that neither optimize nor +judge. +""" +from __future__ import annotations + +import logging + +from rdkit import Chem +from rdkit.Chem import AllChem + +from Auto3D.constants import MIN_ATOM_DISTANCE +from Auto3D.utils.geometry import min_pairwise_distance +from Auto3D.utils.stereo_check import stereo_descriptors_from_3d + +logger = logging.getLogger("auto3d") + +__all__ = ["relieve_clash"] + + +def relieve_clash( + mol: Chem.Mol, + conf_id: int, + min_distance: float = MIN_ATOM_DISTANCE, +) -> bool: + """Optimize a clashing conformer in place and report whether it is usable. + + A conformer is considered "clashing" when its minimum pairwise interatomic + distance is below ``min_distance``. Such conformers are relaxed with MMFF; + when the molecule lacks full MMFF parameters (elements like B, Se or some + Si valences, where ``MMFFOptimizeMolecule`` returns -1 and does nothing), + the function falls back to UFF so the conformer is not discarded for lack + of a force field. + + The force-field relaxation can itself invert a stereocenter or rotate a + double bond. This runs before the enumerated SDF is written, so the + downstream post-optimization stereochemistry check would otherwise read an + already-changed geometry as its own "before" reference and never notice. + Stereochemistry is therefore checked before and after the relaxation, on + this same molecule object, and a conformer whose configuration changed is + rejected here rather than passed downstream. + + Known limitation: the "before" snapshot is read while the conformer is + still in violation of ``min_distance`` — by definition, since that is the + only way execution reaches this branch. CIP perception on a geometry that + is itself clashing is not a trustworthy baseline, unlike the equivalent + check in ``batch_opt/batchopt.py``, whose "before" reading is always + taken from a valid, non-clashing conformer. This matters only when the + branch is actually reached: across roughly 650 conformers sampled from + Auto3D's real ``EmbedMultipleConfs`` output (glucose, cholesterol, a + tripeptide, macrocycles, a cage compound, and molecules with B/Se/ + hypervalent Si), none ever fell below the clash threshold. Under 196 + artificially forced clashes, this guard rejected 96 conformers, and about + 56% of those rejections had a post-relaxation configuration that actually + matched the molecule's true configuration -- spurious rejections caused + by the unreliable baseline rather than a real inversion. The known + improvement is to compare against the molecule's graph-encoded stereo + tags instead of a 3D read of the clashing geometry, but that needs its + own measurement first: RDKit's graph ``AssignStereochemistry`` and + ``AssignStereochemistryFrom3D`` label pseudoasymmetric centers + differently (``r``/``s`` vs ``R``/``S``), which could introduce a + systematic false positive. + + Args: + mol: RDKit molecule holding the conformer. + conf_id: Index of the conformer to check/optimize. + min_distance: Minimum acceptable interatomic distance (Angstroms). + + Returns: + True if the (possibly optimized) conformer's minimum pairwise distance + is >= ``min_distance`` and its stereochemistry survived unchanged; + False if it still clashes or if the relaxation changed its + configuration. + """ + positions = mol.GetConformer(conf_id).GetPositions() + # Closing the dead band: a conformer exactly at the threshold is kept. + if min_pairwise_distance(positions) >= min_distance: + return True + + # Clashing conformer: try MMFF, fall back to UFF when MMFF is unavailable. + before = stereo_descriptors_from_3d(mol, conf_id=conf_id) + if AllChem.MMFFHasAllMoleculeParams(mol): + AllChem.MMFFOptimizeMolecule(mol, confId=conf_id) + else: + AllChem.UFFOptimizeMolecule(mol, confId=conf_id) + + # Clash relief is a force-field relaxation and can invert a center just as + # the neural network optimization can. It runs before the enumerated SDF is + # written, so the post-optimization check downstream would read an already + # inverted geometry as its reference and never notice. Reject the conformer + # here instead; the embedder simply keeps the ones that survive. + if stereo_descriptors_from_3d(mol, conf_id=conf_id) != before: + logger.warning( + "Discarding a conformer whose stereochemistry changed during clash " + "relief." + ) + return False + + positions = mol.GetConformer(conf_id).GetPositions() + return min_pairwise_distance(positions) >= min_distance diff --git a/src/Auto3D/cli/commands/config.py b/src/Auto3D/cli/commands/config.py index f1e4320c..59dd350c 100644 --- a/src/Auto3D/cli/commands/config.py +++ b/src/Auto3D/cli/commands/config.py @@ -65,6 +65,19 @@ } +def _require_config_file_exists(config_file: Path) -> None: + """Raise the standard "config file not found" error if it is missing. + + Both ``execute_config_show`` and ``execute_config_validate`` had this + exact check (message and hint, word for word) inlined separately. + """ + if not config_file.exists(): + raise ConfigurationError( + f"Config file not found: {config_file}", + hint="Run 'auto3d config init' to create one.", + ) + + def generate_commented_yaml(config: dict) -> str: """Generate YAML with helpful comments.""" lines = ["# Auto3D Configuration File", "# Generated by: auto3d config init", ""] @@ -156,11 +169,7 @@ def execute_config_show(config_file: Path | None = None, verbose: int = 0) -> No if config_file is None: config_file = Path("auto3d.yaml") - if not config_file.exists(): - raise ConfigurationError( - f"Config file not found: {config_file}", - hint="Run 'auto3d config init' to create one.", - ) + _require_config_file_exists(config_file) content = config_file.read_text() @@ -192,14 +201,10 @@ def execute_config_validate(config_file: Path, verbose: int = 0) -> None: ``CLIConfig.path``. """ try: - if not config_file.exists(): - # Unreachable from the CLI (Typer's `exists=True` on the argument - # rejects a missing path as a usage error, also exit 2), but this - # function is callable directly. - raise ConfigurationError( - f"Config file not found: {config_file}", - hint="Run 'auto3d config init' to create one.", - ) + # Unreachable from the CLI (Typer's `exists=True` on the argument + # rejects a missing path as a usage error, also exit 2), but this + # function is callable directly. + _require_config_file_exists(config_file) config = load_yaml_config(config_file) diff --git a/src/Auto3D/cli/commands/models.py b/src/Auto3D/cli/commands/models.py index aa83fcbe..98a70887 100644 --- a/src/Auto3D/cli/commands/models.py +++ b/src/Auto3D/cli/commands/models.py @@ -230,8 +230,10 @@ def execute_models_test( ANI, a failed/blocked aimnet registry download, or a broken custom model file -- instead of having them surface deep inside a run. """ + # handle_error is already imported at module level (used identically by + # execute_models_info above); this used to re-import it locally too, a + # dead duplicate of the same binding. from Auto3D.cli.console import print_success - from Auto3D.cli.errors import handle_error try: import time diff --git a/src/Auto3D/cli/commands/run.py b/src/Auto3D/cli/commands/run.py index 5b36fb5e..0e283a34 100644 --- a/src/Auto3D/cli/commands/run.py +++ b/src/Auto3D/cli/commands/run.py @@ -21,7 +21,7 @@ print_failures, print_results_summary, ) -from Auto3D.exceptions import Auto3DError, ConfigurationError +from Auto3D.exceptions import ConfigurationError from Auto3D.utils.logging_config import configure_logging # A partial run (the process completed, but some input molecules produced no @@ -303,10 +303,10 @@ def progress_cb(event: dict) -> None: _exit_if_incomplete(results) except KeyboardInterrupt: - # KeyboardInterrupt is a BaseException, so neither `except - # Auto3DError` nor `except Exception` below ever saw it: Ctrl-C - # printed nothing at all and left the user with no idea how far the - # run had got or whether anything reached disk. + # KeyboardInterrupt is a BaseException, so `except Exception` + # below never saw it: Ctrl-C printed nothing at all and left the + # user with no idea how far the run had got or whether anything + # reached disk. # # Note what this clause is NOT for: typer/core.py already converts # an escaping KeyboardInterrupt into click's Exit(130), so the exit @@ -322,7 +322,11 @@ def progress_cb(event: dict) -> None: batch=display.as_batch_counts() if display is not None else None, elapsed_seconds=time.time() - start_time, ) - except Auto3DError as e: - handle_error(e, verbose=verbose, json_output=json_output) except Exception as e: + # A single clause, not one for Auto3DError and an identical one + # for Exception: handle_error already branches on + # isinstance(error, Auto3DError) internally (cli/errors.py), so a + # separate `except Auto3DError as e: handle_error(...)` ahead of + # this called the exact same function with the exact same + # arguments -- dead structure, not a behavioral distinction. handle_error(e, verbose=verbose, json_output=json_output) diff --git a/src/Auto3D/config.py b/src/Auto3D/config.py index 60683018..189eefef 100644 --- a/src/Auto3D/config.py +++ b/src/Auto3D/config.py @@ -332,7 +332,7 @@ class Auto3DOptions: """Maximum conformers per SMILES. ``None`` derives the count per molecule via - :func:`Auto3D.utils.chemistry.calculate_conformer_count`, which is + :func:`Auto3D.utils.molprops.calculate_conformer_count`, which is ``min(max(1, num_heavy, 2 * 8.481 * num_rotatable ** 1.642), 1000)`` (https://doi.org/10.1021/acs.jctc.0c01213). The rotatable-bond term dominates for anything flexible: glycerol gets **238**, not 5. This docstring used to diff --git a/src/Auto3D/constants.py b/src/Auto3D/constants.py index 1b808aae..5047c3b0 100644 --- a/src/Auto3D/constants.py +++ b/src/Auto3D/constants.py @@ -5,6 +5,12 @@ HARTREE_TO_EV = 27.211386245988 # 1 Hartree in eV EV_TO_KCAL_PER_MOL = 23.060547830619026 # 1 eV in kcal/mol (original Auto3D value) HARTREE_TO_KCAL_PER_MOL = 627.50947337481 # 1 Hartree in kcal/mol +# 1 eV in Hartree. Computed from HARTREE_TO_EV (not an independent literal), so +# the two can never drift apart. Previously defined independently -- misspelled +# as `ev2hatree`, and via the identical expression `1 / hartree2ev` -- in both +# ASE/thermo.py and SPE.py (M62); this is the single, correctly spelled owner +# both now import instead of recomputing their own copy. +EV_TO_HARTREE = 1.0 / HARTREE_TO_EV # Geometry thresholds MIN_ATOM_DISTANCE = 0.9 # Å, minimum allowed interatomic distance diff --git a/src/Auto3D/isomers/parallel_embed.py b/src/Auto3D/embedding.py similarity index 89% rename from src/Auto3D/isomers/parallel_embed.py rename to src/Auto3D/embedding.py index 65607775..a5fbf232 100644 --- a/src/Auto3D/isomers/parallel_embed.py +++ b/src/Auto3D/embedding.py @@ -1,5 +1,15 @@ -# src/Auto3D/isomers/parallel_embed.py -"""Parallel conformer embedding using multiprocessing.""" +# src/Auto3D/embedding.py +"""Parallel conformer embedding using multiprocessing. + +Lives at the top level rather than under ``Auto3D.isomers`` because +``isomer_engine`` is its only caller and ``isomers`` is the package that +*wraps* ``isomer_engine``: with this module inside ``isomers``, the two +packages imported each other (``isomers.factory``/the adapters reached into +``isomer_engine``, and ``isomer_engine._run_parallel_embedding`` reached back +into ``isomers.parallel_embed``), a cycle that only stayed latent because +every edge of it was a function-scope import. Moving this module out is what +removes the cycle rather than deferring it. +""" from __future__ import annotations from collections.abc import Iterator @@ -9,9 +19,10 @@ from rdkit import Chem from rdkit.Chem import AllChem +from Auto3D.clash_relief import relieve_clash from Auto3D.constants import CONFORMER_RANDOM_SEED -from Auto3D.utils.chemistry import calculate_conformer_count, relieve_clash from Auto3D.utils.logging_config import get_logger +from Auto3D.utils.molprops import calculate_conformer_count logger = get_logger(__name__) diff --git a/src/Auto3D/filtering.py b/src/Auto3D/filtering.py index 047a6874..b4af3b98 100644 --- a/src/Auto3D/filtering.py +++ b/src/Auto3D/filtering.py @@ -10,7 +10,7 @@ DEFAULT_ENERGY_CLUSTER_WINDOW, DEFAULT_RMSD_THRESHOLD, ) -from Auto3D.utils.chemistry import check_connectivity +from Auto3D.utils.connectivity import check_connectivity from Auto3D.utils.convergence import converged_or_unfiltered from Auto3D.utils.energy import e_tot_ev, try_e_tot_ev from Auto3D.utils.stereo_check import species_key, stereo_preserved @@ -25,7 +25,7 @@ def filter_unique_optimized( Sorts by energy and only RMSD-compares molecules close enough in energy to be duplicates at all, which avoids the O(n^2) comparisons of the legacy - ``utils.chemistry.filter_unique`` **without changing which molecules + :func:`filter_unique` below **without changing which molecules survive** -- the partitioning is chosen so that no duplicate pair can be separated by it. See the comment on the split rule below for why that holds; it did not hold before 4.0.0. @@ -200,8 +200,111 @@ def _mol_energy(mol: Chem.Mol) -> float | None: """Return a mol's optimized energy in eV, or None. ``E_tot`` is stored in Hartree; the conversion lives in - ``Auto3D.utils.energy`` so this module and ``ranking``/``utils.chemistry`` - cannot drift apart on it. ``None`` signals "no usable energy" so callers - fall back to RMSD-only comparison. + ``Auto3D.utils.energy`` so this module and ``ranking`` cannot drift apart on + it. ``None`` signals "no usable energy" so callers fall back to RMSD-only + comparison. """ return try_e_tot_ev(mol) + + +def filter_unique(mols: list[Chem.Mol], crit: float = DEFAULT_RMSD_THRESHOLD) -> list[Chem.Mol]: + """Remove structures that are very similar and remove unconverged structures. + + The legacy all-pairs filter, kept beside :func:`filter_unique_optimized` + (which supersedes it) because ``ConformerRanker(use_optimized_filtering=False)`` + still selects it and because it is the oracle the optimized filter is + compared against. + + This function filters a list of molecules to keep only unique, converged structures. + It first removes unconverged structures and those with invalid connectivity, + then removes similar structures based on RMSD comparison. + + Args: + mols: List of RDKit molecule objects, optionally carrying a 'Converged' + property. A record whose 'Converged' is explicitly false is + dropped; a record without the property is kept (not filtered on + convergence). Records marked 'Stereo_changed' are excluded. + crit: RMSD threshold for considering two structures as identical. + Structures with RMSD below this value are considered duplicates. + Defaults to DEFAULT_RMSD_THRESHOLD (0.3 Angstroms). + + Returns: + List of unique, converged molecules with valid connectivity. + + Example: + >>> from rdkit import Chem + >>> from rdkit.Chem import AllChem + >>> mol = Chem.MolFromSmiles("CCO") + >>> mol = Chem.AddHs(mol) + >>> AllChem.EmbedMolecule(mol, randomSeed=42) + 0 + >>> mol.SetProp("Converged", "true") + >>> filter_unique([mol], crit=0.3) # Returns list with 1 molecule + [...] + """ + # Remove structures that explicitly failed to converge. A record with no + # 'Converged' property is NOT filtered on convergence -- see + # Auto3D.utils.convergence for why absence is not failure. + mols_: list[Chem.Mol] = [] + for mol in mols: + convergence_flag = converged_or_unfiltered(mol) + has_valid_bonds = check_connectivity(mol) + if convergence_flag and has_valid_bonds and stereo_preserved(mol): + mols_.append(mol) + mols = mols_ + + # Remove similar structures. Strip Hs once per molecule (O(n)) instead of on + # both sides of every comparison (O(n^2)); GetBestRMS on no-H forms is + # symmetric so results are unchanged. The ORIGINAL (H-explicit) mols are + # returned; no-H forms are comparison-only. + # + # Heavy-atom RMSD alone collapses conformers that differ only in an O-H / N-H + # rotor orientation. Guard with an energy check: a pair counts as duplicate + # only when the RMSD is below ``crit`` AND the energies agree within + # DEFAULT_DUPLICATE_ENERGY_TOL (eV; 'E_tot' is stored in Hartree and is + # converted on read by Auto3D.utils.energy). Mols without a usable 'E_tot' + # fall back to RMSD-only (energy guard cannot apply). + unique_mols: list[Chem.Mol] = [] + unique_noH: list[Chem.Mol] = [] + unique_energies: list[float | None] = [] + unique_species: list[str] = [] + for mol_i in mols: + mol_i_noH = Chem.RemoveHs(mol_i) + # E_tot is stored in Hartree; DEFAULT_DUPLICATE_ENERGY_TOL is in eV. + e_i: float | None = try_e_tot_ev(mol_i) + species_i = species_key(mol_i) + unique = True + for mol_j_noH, e_j, species_j in zip( + unique_noH, unique_energies, unique_species, strict=True + ): + # Two different compounds are never duplicates of each other, however + # close their geometries. All stereoisomers of one input share a + # ranking group, and two ring diastereomers can sit below the default + # 0.3 A threshold, so without this the RMSD test could delete one of + # them (Auto3D.utils.stereo_check.species_key). Checked before the + # RMSD call it makes unnecessary. + if species_i != species_j: + continue + try: + # temporary bug fix for https://github.com/rdkit/rdkit/issues/6826 + # removing Hs speeds up the calculation + rmsd = rdMolAlign.GetBestRMS(mol_i_noH, mol_j_noH) + except RuntimeError: + # Incomparable pair: treat as distinct (not a duplicate) so the + # conformer is kept. Using 0 would make it look like a perfect + # duplicate and drop a genuinely distinct structure. + rmsd = float("inf") + energy_close = ( + e_i is None + or e_j is None + or abs(e_i - e_j) < DEFAULT_DUPLICATE_ENERGY_TOL + ) + if rmsd < crit and energy_close: + unique = False + break + if unique: + unique_mols.append(mol_i) + unique_noH.append(mol_i_noH) + unique_energies.append(e_i) + unique_species.append(species_i) + return unique_mols diff --git a/src/Auto3D/id_mapping.py b/src/Auto3D/id_mapping.py new file mode 100644 index 00000000..7b2c45e4 --- /dev/null +++ b/src/Auto3D/id_mapping.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python +"""Auto3D's numeric molecule-ID encoding, and its inverse. + +Pipeline policy rather than file I/O: the run replaces every user ID with a +dense integer index so that the chunking, isomer and optimizer stages can key +on something short and collision-free, then restores the original IDs on the +way out. Both halves must agree on the index space and on how a ``@tautN`` +suffix survives it, which is why they live together and above +``Auto3D.utils`` -- ``utils`` is a leaf of generic helpers, and this is a +decision about how an Auto3D run is shaped. +""" +from __future__ import annotations + +import os +from pathlib import Path + +from rdkit import Chem + +from Auto3D.exceptions import ConfigurationError, InputValidationError +from Auto3D.utils.logging_config import get_logger +from Auto3D.utils.smi_io import iter_smi_records + +logger = get_logger(__name__) + + +def encode_ids( + path: str, out_dir: str | os.PathLike[str] | None = None +) -> tuple[str, dict[str, int]]: + """Encode molecule IDs to numeric indices. + + For a .smi or .sdf file, replaces all molecule IDs with sequential + integer indices and returns a mapping from original IDs to indices. + + The encoded file is named ``_encoded.``. That name is derived + from the input, so it can collide with a file the user already owns: + ``mols_encoded.smi`` sitting beside ``mols.smi`` is a perfectly ordinary + thing for a user to have, and this function used to overwrite it without + a word (``WorkflowOrchestrator`` then ``unlink()``ed it at the end of the + run, so the file was destroyed twice over). Two things prevent that now: + ``out_dir`` lets the caller redirect the encoded file somewhere it owns + -- ``WorkflowOrchestrator`` passes its freshly created job directory -- + and the collision check below refuses to write over an existing file for + every caller, including ones that take the default location. + + Args: + path: Path to the input .smi or .sdf file. + out_dir: Directory to write the encoded file into. Defaults to the + input file's own directory. + + Returns: + Tuple containing: + - Path to the new file with encoded IDs (adds '_encoded' suffix) + - Dictionary mapping original IDs to their numeric indices + + Raises: + ValueError: If the input file is neither .smi nor .sdf format. + ConfigurationError: If a file already exists at the encoded path. + InputValidationError: If a molecule has a missing/blank ID or a + duplicate ID is encountered. + + Example: + >>> new_path, mapping = encode_ids("molecules.smi") + >>> mapping + {'mol_A': 0, 'mol_B': 1, 'mol_C': 2} + """ + path_obj = Path(path).resolve() + extension = path_obj.suffix[1:] + # Checked up front rather than in a trailing `else`: the collision check + # below must not be the thing that reports an unsupported extension. + if extension not in ("smi", "sdf"): + raise ValueError("The input file should be either smi or sdf") + + directory = Path(out_dir) if out_dir is not None else path_obj.parent + new_path = directory / f"{path_obj.stem}_encoded.{extension}" + if new_path.exists(): + raise ConfigurationError( + f"encode_ids would overwrite the existing file {new_path}. " + "Auto3D writes its encoded copy of the input there; move or " + "rename that file, or pass out_dir to write the encoded copy " + "somewhere else." + ) + + if extension == "smi": + new_data: list[str] = [] + mapping: dict[str, int] = {} + # iter_smi_records raises InputValidationError on a <2-token line + # (on_malformed="raise"). Duplicate-id detection stays here because the + # helper does not dedup. Index by a dense record counter, not the file + # line number: blank/skipped lines would otherwise leave gaps in the + # index space, which is inconsistent with the dense positions the chunk + # manager assumes downstream. The original file line_no is still used in + # the error message so it points at the real offending line. + for i, (line_no, smi, id) in enumerate( + iter_smi_records(path, on_malformed="raise") + ): + if id in mapping: + raise InputValidationError( + f"Duplicate molecule ID {id!r} on line {line_no}. " + "IDs must be unique." + ) + mapping[id] = i + new_data.append(f"{smi} {i}\n") + with open(new_path, "w") as f: + for line in new_data: + f.write(line) + return str(new_path), mapping + + else: # "sdf" -- the only remaining possibility, checked above + suppl = Chem.SDMolSupplier(path, removeHs=False) + mapping = {} + with Chem.SDWriter(str(new_path)) as w: + for i, mol in enumerate(suppl): + if mol is None: + logger.warning(f"Skipping molecule at index {i}: failed to parse") + continue + id = mol.GetProp("_Name").strip() + if not id: + raise InputValidationError( + f"Molecule at index {i} has a missing or blank name." + ) + if id in mapping: + raise InputValidationError( + f"Duplicate molecule name {id!r} at index {i}. " + "Names must be unique." + ) + mapping[id] = i + mol.SetProp("_Name", str(i)) + w.write(mol) + return str(new_path), mapping + + +def decode_ids(path: str, mapping: dict[str, int]) -> str: + """Decode numeric IDs back to original molecule IDs. + + For an SDF file with numeric IDs, restores the original IDs using + the provided mapping dictionary. + + Args: + path: Path to the input SDF file with encoded (numeric) IDs. + mapping: Dictionary mapping original IDs to their numeric indices + (as returned by encode_ids). + + Returns: + Path to the new SDF file with decoded IDs (adds '_out' suffix). + + Example: + >>> mapping = {'mol_A': 0, 'mol_B': 1} + >>> output_path = decode_ids("encoded_3d.sdf", mapping) + """ + # Invert the mapping: index -> original_id + inverse_mapping = {v: k for k, v in mapping.items()} + path_obj = Path(path).resolve() + extension = path_obj.suffix[1:] + # Reconstruct base name: remove last two underscore-separated parts + stem_parts = path_obj.stem.split("_")[:-2] + new_stem = "_".join(stem_parts) + "_out" + new_path = path_obj.parent / f"{new_stem}.{extension}" + + suppl = Chem.SDMolSupplier(path, removeHs=False) + with Chem.SDWriter(str(new_path)) as w: + for i, mol in enumerate(suppl): + if mol is None: + logger.warning("Skipping molecule at index %d: failed to parse", i) + continue + name = mol.GetProp("_Name").strip() + if "@taut" in name: + components = name.split("@taut") + new_name = ( + inverse_mapping[int(components[0])] + "@taut" + "".join(components[1:]) + ) + else: + new_name = inverse_mapping[int(name)] + mol.SetProp("_Name", new_name) + + id = "_".join(mol.GetProp("ID").strip().split("_")[1:]) + new_id = new_name + "_" + id + mol.SetProp("ID", new_id) + + w.write(mol) + return str(new_path) diff --git a/src/Auto3D/isomer_engine.py b/src/Auto3D/isomer_engine.py index de8ba408..aaa14aaa 100644 --- a/src/Auto3D/isomer_engine.py +++ b/src/Auto3D/isomer_engine.py @@ -19,9 +19,10 @@ from rdkit.Chem.MolStandardize import rdMolStandardize from tqdm import tqdm +from Auto3D.clash_relief import relieve_clash from Auto3D.constants import CONFORMER_RANDOM_SEED, MAX_STEREOISOMERS -from Auto3D.utils.chemistry import calculate_conformer_count, relieve_clash -from Auto3D.utils.file_ops import ( +from Auto3D.utils.molprops import calculate_conformer_count +from Auto3D.utils.smi_io import ( combine_smi, hash_enumerated_smi_IDs, iter_smi_records, @@ -399,7 +400,9 @@ def _run_parallel_embedding( self, smi_name_tuples: list[tuple[str, str]] ) -> None: """Run parallel conformer embedding using ProcessPoolExecutor.""" - from Auto3D.isomers.parallel_embed import embed_conformers_parallel + # Function-scope on purpose: tests patch the attribute on + # ``Auto3D.embedding`` and rely on this lookup re-reading it. + from Auto3D.embedding import embed_conformers_parallel logger.info(f"Using parallel embedding with {self.parallel_workers} workers...") @@ -428,7 +431,7 @@ class RDKitSdfIsomer: isomer, so this path's own output has one consistent shape to parse -- the same shape the SMILES path emits, with or without ``enumerate_isomers``. The isomer component is what - :func:`Auto3D.utils.file_ops.decode_ids` relies on to rebuild + :func:`Auto3D.id_mapping.decode_ids` relies on to rebuild ``__`` IDs after the pipeline's numeric-ID encoding step; :class:`~Auto3D.ranking.ConformerRanker` groups on the leading component only, so it is unaffected by the isomer index. @@ -601,7 +604,7 @@ def oe_isomer( The OpenEye toolkit's application-options machinery writes ``oeomega_*`` and ``flipper_*`` logfiles into the **process working directory**, which for an ordinary ``cd ~/project && auto3d run mols.smi`` is the user's own - directory. ``utils.file_ops.housekeeping`` used to sweep those names out + directory. ``job_layout.housekeeping`` used to sweep those names out of the cwd and into the run's ``verbose`` folder, which is tarred and then deleted -- so a user file named ``oeomega_settings.txt`` in the cwd was destroyed by an ordinary run. The fix is on this side rather than on the diff --git a/src/Auto3D/isomers/__init__.py b/src/Auto3D/isomers/__init__.py index 18e7e7af..57d9e502 100644 --- a/src/Auto3D/isomers/__init__.py +++ b/src/Auto3D/isomers/__init__.py @@ -1,34 +1,24 @@ """Isomer enumeration engines with Strategy Pattern. -This module provides a unified interface for different isomer enumeration +This package provides a unified interface for different isomer enumeration backends (RDKit, OpenEye Omega) through the Strategy Pattern. +``IsomerEngineFactory`` is the one name re-exported here, because +``docs/source/api.rst`` documents it at this package path +(``Auto3D.isomers.IsomerEngineFactory``) and that path is the public one. Every +other name in this package -- ``create_isomer_engine``, +``create_tautomer_engine``, and the ``IsomerEngine``/``TautomerEngine`` +protocols -- is imported from the module that defines it +(``Auto3D.isomers.factory``, ``Auto3D.isomers.base``), so that no name here has +two supported spellings. + Example: >>> from Auto3D.isomers import IsomerEngineFactory >>> engine = IsomerEngineFactory.create("rdkit", input_path="input.smi", ...) >>> engine.run() - - # Or using the convenience function - >>> from Auto3D.isomers import create_isomer_engine - >>> engine = create_isomer_engine("rdkit", input_path="input.smi", ...) """ from __future__ import annotations -from Auto3D.isomers.base import BaseIsomerEngine, IsomerEngine, TautomerEngine -from Auto3D.isomers.factory import ( - IsomerEngineFactory, - create_isomer_engine, - create_tautomer_engine, -) +from Auto3D.isomers.factory import IsomerEngineFactory -__all__ = [ - # Protocols and base classes - "IsomerEngine", - "TautomerEngine", - "BaseIsomerEngine", - # Factory - "IsomerEngineFactory", - # Convenience functions - "create_isomer_engine", - "create_tautomer_engine", -] +__all__ = ["IsomerEngineFactory"] diff --git a/src/Auto3D/isomers/base.py b/src/Auto3D/isomers/base.py index df6a8608..75e98553 100644 --- a/src/Auto3D/isomers/base.py +++ b/src/Auto3D/isomers/base.py @@ -1,7 +1,7 @@ """Base protocols for isomer enumeration engines.""" from __future__ import annotations -from abc import ABC, abstractmethod +from abc import abstractmethod from typing import Protocol, runtime_checkable @@ -44,42 +44,3 @@ class TautomerEngine(Protocol): def run(self) -> None: """Execute tautomer enumeration.""" ... - - -class BaseIsomerEngine(ABC): - """Abstract base class for isomer engines. - - Provides common functionality for all isomer engine implementations. - """ - - def __init__( - self, - input_path: str, - output_path: str, - max_confs: int | None = None, - threshold: float = 0.3, - n_jobs: int = 4, - ) -> None: - """Initialize the isomer engine. - - Args: - input_path: Path to input file (SMI or SDF). - output_path: Path for output SDF file. - max_confs: Maximum conformers per molecule. None for dynamic. - threshold: RMSD threshold for duplicate removal (Å). - n_jobs: Number of parallel jobs for conformer generation. - """ - self.input_path = input_path - self.output_path = output_path - self.max_confs = max_confs - self.threshold = threshold - self.n_jobs = n_jobs - - @abstractmethod - def run(self) -> str: - """Execute isomer enumeration. - - Returns: - Path to the output SDF file. - """ - ... diff --git a/src/Auto3D/isomers/factory.py b/src/Auto3D/isomers/factory.py index c71a0543..04a78381 100644 --- a/src/Auto3D/isomers/factory.py +++ b/src/Auto3D/isomers/factory.py @@ -1,12 +1,50 @@ """Factory classes and functions for creating isomer engines.""" from __future__ import annotations -from typing import TYPE_CHECKING +from collections.abc import Callable + +from Auto3D.isomer_engine import ( + RDKitIsomer, + RDKitSdfIsomer, + oe_isomer, +) +from Auto3D.isomer_engine import ( + TautomerEngine as _RDKitOrOmegaTautomerEngine, +) +from Auto3D.isomers.base import IsomerEngine, TautomerEngine + +#: Engine names :meth:`IsomerEngineFactory.create` accepts. ``rdkit`` is also +#: reachable as ``rdkit_sdf`` via the ``input_format="sdf"`` auto-selection +#: below. +_ENGINE_TYPES = ("rdkit", "rdkit_sdf", "omega") + + +class _DeferredIsomerEngine: + """The ``.run() -> str`` object :meth:`IsomerEngineFactory.create` returns. + + Holds nothing but the zero-argument callable that builds the concrete + engine and runs it. Construction stays deferred to ``run()`` on purpose and + is not an implementation detail: ``RDKitIsomer.__init__`` creates its + ``rdk_tmp`` working directory, so building the engine at ``create()`` time + would move a filesystem side effect (and its ``FileExistsError``) earlier + than every caller was written for. + + This one class replaces three adapter classes and an abstract base that + existed only to copy the same eight-to-thirteen keyword arguments twice -- + once into an adapter's ``__init__``, once out of it into the real engine. + ``create`` below now maps them once, at the single site that knows what + each engine takes. + """ + + __slots__ = ("_build_and_run", "output_path") -from Auto3D.isomers.base import BaseIsomerEngine, IsomerEngine, TautomerEngine + def __init__(self, build_and_run: Callable[[], str], output_path: str) -> None: + self._build_and_run = build_and_run + self.output_path = output_path -if TYPE_CHECKING: - pass + def run(self) -> str: + """Build the concrete engine, run it, and return its output path.""" + return self._build_and_run() class IsomerEngineFactory: @@ -26,27 +64,6 @@ class IsomerEngineFactory: >>> output = engine.run() """ - # Registry of available adapters - _adapters: dict[str, type[BaseIsomerEngine]] = {} - - @classmethod - def _ensure_registered(cls) -> None: - """Lazily register adapters on first use.""" - if cls._adapters: - return - - from Auto3D.isomers.omega_adapter import OmegaIsomerAdapter - from Auto3D.isomers.rdkit_adapters import ( - RDKitIsomerAdapter, - RDKitSdfIsomerAdapter, - ) - - cls._adapters = { - "rdkit": RDKitIsomerAdapter, - "rdkit_sdf": RDKitSdfIsomerAdapter, - "omega": OmegaIsomerAdapter, - } - @classmethod def available_engines(cls) -> list[str]: """Return list of available engine types. @@ -54,8 +71,7 @@ def available_engines(cls) -> list[str]: Returns: List of supported engine type names. """ - cls._ensure_registered() - return list(cls._adapters.keys()) + return list(_ENGINE_TYPES) @classmethod def create( @@ -77,7 +93,7 @@ def create( use_parallel_embedding: bool = False, parallel_embedding_threshold: int = 10, parallel_workers: int = 4, - ) -> BaseIsomerEngine: + ) -> IsomerEngine: """Create an isomer engine adapter. Args: @@ -100,68 +116,77 @@ def create( parallel_workers: Number of worker processes for parallel embedding. Returns: - Configured isomer engine adapter instance. + Configured isomer engine, whose ``run()`` builds and drives the + concrete engine and returns the output path. Raises: ValueError: If engine_type is not recognized. """ - cls._ensure_registered() engine_type = engine_type.lower() # Auto-select rdkit_sdf for SDF input when rdkit is requested if engine_type == "rdkit" and input_format.lower() == "sdf": engine_type = "rdkit_sdf" - if engine_type not in cls._adapters: - available = ", ".join(f"'{e}'" for e in cls._adapters.keys()) + if engine_type not in _ENGINE_TYPES: + available = ", ".join(f"'{e}'" for e in _ENGINE_TYPES) raise ValueError( f"Unknown isomer engine type: '{engine_type}'. " f"Supported types: {available}" ) - adapter_class = cls._adapters[engine_type] - - # Create adapter with appropriate parameters + # The one kwarg-mapping site. Each branch names exactly the arguments + # its engine takes, so an argument no engine reads cannot survive here + # unnoticed the way it could when three adapters each held a partial + # copy of the same signature. if engine_type == "rdkit": - return adapter_class( - input_path=input_path, - output_path=output_path, - smiles_enumerated=smiles_enumerated, - smiles_reduced=smiles_reduced, - smiles_hashed=smiles_hashed, - job_dir=job_dir, - max_confs=max_confs, - threshold=threshold, - n_jobs=n_jobs, - enumerate_isomers=enumerate_isomers, - use_parallel_embedding=use_parallel_embedding, - parallel_embedding_threshold=parallel_embedding_threshold, - parallel_workers=parallel_workers, - ) - elif engine_type == "rdkit_sdf": - return adapter_class( - input_path=input_path, - output_path=output_path, - max_confs=max_confs, - threshold=threshold, - n_jobs=n_jobs, - enumerate_isomers=enumerate_isomers, - ) - elif engine_type == "omega": - return adapter_class( - mode=mode, - input_path=input_path, - smiles_enumerated=smiles_enumerated, - smiles_reduced=smiles_reduced, - smiles_hashed=smiles_hashed, - output_path=output_path, - max_confs=max_confs, - threshold=threshold, - enumerate_isomers=enumerate_isomers, - ) + def build_and_run() -> str: + return RDKitIsomer( + smi=input_path, + smiles_enumerated=smiles_enumerated, + smiles_enumerated_reduced=smiles_reduced, + smiles_hashed=smiles_hashed, + enumerated_sdf=output_path, + job_name=job_dir, + max_confs=max_confs, + threshold=threshold, + np=n_jobs, + flipper=enumerate_isomers, + use_parallel_embedding=use_parallel_embedding, + parallel_embedding_threshold=parallel_embedding_threshold, + parallel_workers=parallel_workers, + ).run() - # Should not reach here, but for type safety - raise ValueError(f"Unknown engine type: {engine_type}") + elif engine_type == "rdkit_sdf": + def build_and_run() -> str: + return RDKitSdfIsomer( + sdf=input_path, + enumerated_sdf=output_path, + max_confs=max_confs, + threshold=threshold, + np=n_jobs, + flipper=enumerate_isomers, + ).run() + + else: # "omega" -- the only remaining possibility, checked above + def build_and_run() -> str: + # oe_isomer is a function returning 0, not an engine object, so + # this is the one branch that still has an adapting step: run it + # and report the path it wrote. + oe_isomer( + mode=mode, + input_f=input_path, + smiles_enumerated=smiles_enumerated, + smiles_reduced=smiles_reduced, + smiles_hashed=smiles_hashed, + output=output_path, + max_confs=max_confs, + threshold=threshold, + flipper=enumerate_isomers, + ) + return output_path + + return _DeferredIsomerEngine(build_and_run, output_path) def create_isomer_engine( @@ -270,9 +295,7 @@ def create_tautomer_engine( engine_type = engine_type.lower() if engine_type in ("rdkit", "oechem"): - from Auto3D.isomer_engine import TautomerEngine as TautEngine - - return TautEngine( + return _RDKitOrOmegaTautomerEngine( mode=engine_type, input_f=input_path, out=output_path, diff --git a/src/Auto3D/isomers/omega_adapter.py b/src/Auto3D/isomers/omega_adapter.py deleted file mode 100644 index 9b87ae24..00000000 --- a/src/Auto3D/isomers/omega_adapter.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Adapter for OpenEye Omega isomer engine.""" -from __future__ import annotations - -from Auto3D.isomers.base import BaseIsomerEngine - - -class OmegaIsomerAdapter(BaseIsomerEngine): - """Adapter wrapping the OpenEye Omega isomer generation. - - This adapter provides a consistent interface for the Omega engine, - wrapping the oe_isomer function from isomer_engine.py. - """ - - def __init__( - self, - mode: str, - input_path: str, - smiles_enumerated: str, - smiles_reduced: str, - smiles_hashed: str, - output_path: str, - max_confs: int | None = None, - threshold: float = 0.3, - enumerate_isomers: bool = True, - ) -> None: - """Initialize the Omega adapter. - - Args: - mode: Omega mode ('classic', 'macrocycle', 'dense', 'pose', 'rocs', 'fast_rocs'). - input_path: Path to input SMI or SDF file. - smiles_enumerated: Path for enumerated stereoisomers. - smiles_reduced: Path for reduced isomers (no enantiomers). - smiles_hashed: Path for hashed SMILES IDs. - output_path: Path for output SDF file. - max_confs: Maximum conformers per molecule. - threshold: RMSD threshold for duplicate removal. - enumerate_isomers: Whether to enumerate stereoisomers. - """ - super().__init__( - input_path=input_path, - output_path=output_path, - max_confs=max_confs, - threshold=threshold, - ) - self.mode = mode - self.smiles_enumerated = smiles_enumerated - self.smiles_reduced = smiles_reduced - self.smiles_hashed = smiles_hashed - self.enumerate_isomers = enumerate_isomers - - def run(self) -> str: - """Execute Omega isomer generation. - - Returns: - Path to the output SDF file. - """ - from Auto3D.isomer_engine import oe_isomer - - oe_isomer( - mode=self.mode, - input_f=self.input_path, - smiles_enumerated=self.smiles_enumerated, - smiles_reduced=self.smiles_reduced, - smiles_hashed=self.smiles_hashed, - output=self.output_path, - max_confs=self.max_confs, - threshold=self.threshold, - flipper=self.enumerate_isomers, - ) - - return self.output_path diff --git a/src/Auto3D/isomers/rdkit_adapters.py b/src/Auto3D/isomers/rdkit_adapters.py deleted file mode 100644 index 32d0203d..00000000 --- a/src/Auto3D/isomers/rdkit_adapters.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Adapters for RDKit isomer engines.""" -from __future__ import annotations - -from Auto3D.isomers.base import BaseIsomerEngine - - -class RDKitIsomerAdapter(BaseIsomerEngine): - """Adapter wrapping the RDKitIsomer class. - - This adapter provides a consistent interface for the RDKit SMI-based - isomer generation, following the same pattern as OmegaIsomerAdapter. - """ - - def __init__( - self, - input_path: str, - output_path: str, - smiles_enumerated: str, - smiles_reduced: str, - smiles_hashed: str, - job_dir: str, - max_confs: int | None = None, - threshold: float = 0.3, - n_jobs: int = 4, - enumerate_isomers: bool = True, - use_parallel_embedding: bool = False, - parallel_embedding_threshold: int = 10, - parallel_workers: int = 4, - ) -> None: - """Initialize the RDKit SMI isomer adapter. - - Args: - input_path: Path to input SMI file. - output_path: Path for output SDF file. - smiles_enumerated: Path for enumerated stereoisomers. - smiles_reduced: Path for reduced isomers (no enantiomers). - smiles_hashed: Path for hashed SMILES IDs. - job_dir: Working directory for temporary files. - max_confs: Maximum conformers per molecule. - threshold: RMSD threshold for duplicate removal. - n_jobs: Number of CPU threads for conformer generation. - enumerate_isomers: Whether to enumerate R/S and cis/trans isomers. - use_parallel_embedding: Whether to use parallel conformer embedding. - parallel_embedding_threshold: Minimum molecules for parallel embedding. - parallel_workers: Number of worker processes for parallel embedding. - """ - super().__init__( - input_path=input_path, - output_path=output_path, - max_confs=max_confs, - threshold=threshold, - n_jobs=n_jobs, - ) - self.smiles_enumerated = smiles_enumerated - self.smiles_reduced = smiles_reduced - self.smiles_hashed = smiles_hashed - self.job_dir = job_dir - self.enumerate_isomers = enumerate_isomers - self.use_parallel_embedding = use_parallel_embedding - self.parallel_embedding_threshold = parallel_embedding_threshold - self.parallel_workers = parallel_workers - - def run(self) -> str: - """Execute RDKit SMI-based isomer generation. - - Returns: - Path to the output SDF file. - """ - from Auto3D.isomer_engine import RDKitIsomer - - engine = RDKitIsomer( - smi=self.input_path, - smiles_enumerated=self.smiles_enumerated, - smiles_enumerated_reduced=self.smiles_reduced, - smiles_hashed=self.smiles_hashed, - enumerated_sdf=self.output_path, - job_name=self.job_dir, - max_confs=self.max_confs, - threshold=self.threshold, - np=self.n_jobs, - flipper=self.enumerate_isomers, - use_parallel_embedding=self.use_parallel_embedding, - parallel_embedding_threshold=self.parallel_embedding_threshold, - parallel_workers=self.parallel_workers, - ) - return engine.run() - - -class RDKitSdfIsomerAdapter(BaseIsomerEngine): - """Adapter wrapping the RDKitSdfIsomer class. - - This adapter provides a consistent interface for the RDKit SDF-based - conformer generation, following the same pattern as OmegaIsomerAdapter. - """ - - def __init__( - self, - input_path: str, - output_path: str, - max_confs: int | None = None, - threshold: float = 0.3, - n_jobs: int = 4, - enumerate_isomers: bool = True, - ) -> None: - """Initialize the RDKit SDF isomer adapter. - - Args: - input_path: Path to input SDF file. - output_path: Path for output SDF file. - max_confs: Maximum conformers per stereoisomer. - threshold: RMSD threshold for duplicate removal. - n_jobs: Number of CPU threads for conformer generation. - enumerate_isomers: Whether to enumerate unspecified stereocenters. - """ - super().__init__( - input_path=input_path, - output_path=output_path, - max_confs=max_confs, - threshold=threshold, - n_jobs=n_jobs, - ) - self.enumerate_isomers = enumerate_isomers - - def run(self) -> str: - """Execute RDKit SDF-based conformer generation. - - Returns: - Path to the output SDF file. - """ - from Auto3D.isomer_engine import RDKitSdfIsomer - - engine = RDKitSdfIsomer( - sdf=self.input_path, - enumerated_sdf=self.output_path, - max_confs=self.max_confs, - threshold=self.threshold, - np=self.n_jobs, - flipper=self.enumerate_isomers, - ) - return engine.run() diff --git a/src/Auto3D/job_layout.py b/src/Auto3D/job_layout.py new file mode 100644 index 00000000..f2460d0d --- /dev/null +++ b/src/Auto3D/job_layout.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python +"""Where a run puts its files, and what happens to them afterwards. + +Pipeline policy, not generic file I/O: :func:`create_chunk_meta_names` is the +single place that decides every intermediate file's name inside a chunk's job +directory, and :func:`housekeeping` is what sweeps those same files into the +folder the caller then tars and deletes. The two are a matched pair -- the +sweep is only safe because the names it collects were all minted inside a +directory Auto3D created -- so they live together, above ``Auto3D.utils``. +""" +from __future__ import annotations + +import shutil +from pathlib import Path + +from Auto3D.utils.logging_config import get_logger + +logger = get_logger(__name__) + + +def create_chunk_meta_names(path: str, dir: str) -> dict[str, str]: + """Create output file names based on chunk input path and directory. + + Generates a dictionary of standardized file paths for all intermediate + and output files used in the Auto3D workflow. + + Args: + path: Chunk input .smi file path. + dir: Chunk job folder path. + + Returns: + Dictionary mapping meta names to file paths with the following keys: + - output: Final 3D structure output file + - optimized_og: Original optimized structures + - output_taut: Tautomer SMILES output + - smiles_enumerated: Enumerated SMILES file + - smiles_reduced: Reduced enumerated SMILES file + - smiles_hashed: Hashed enumerated SMILES file + - enumerated_sdf: Enumerated SDF file + - sorted_sdf: Sorted SDF file + - housekeeping_folder: Verbose output folder + - path: Original input path + - dir: Job directory + + Example: + >>> meta = create_chunk_meta_names("chunk1.smi", "/tmp/job") + >>> meta["output"] + '/tmp/job/chunk1_3d.sdf' + """ + dct: dict[str, str] = {} + dir_path = Path(dir) + stem = Path(path).stem + + output = str(dir_path / f"{stem}_3d.sdf") + optimized_og = str(dir_path / f"{stem}_3d0.sdf") + output_taut = str(dir_path / "smi_taut.smi") + smiles_enumerated = str(dir_path / "smiles_enumerated.smi") + smiles_reduced = str(dir_path / "smiles_enumerated_reduced.smi") + smiles_hashed = str(dir_path / "smiles_enumerated_hashed.smi") + enumerated_sdf = str(dir_path / "smiles_enumerated.sdf") + sorted_sdf = str(dir_path / "enumerated_sorted.sdf") + housekeeping_folder = str(dir_path / "verbose") + + dct["output"] = output + dct["optimized_og"] = optimized_og + dct["output_taut"] = output_taut + dct["smiles_enumerated"] = smiles_enumerated + dct["smiles_reduced"] = smiles_reduced + dct["smiles_hashed"] = smiles_hashed + dct["enumerated_sdf"] = enumerated_sdf + dct["sorted_sdf"] = sorted_sdf + dct["housekeeping_folder"] = housekeeping_folder + dct["path"] = path + dct["dir"] = dir + return dct + + +def housekeeping(job_name: str, folder: str, optimized_structures: str) -> None: + """Move this job directory's metadata files into a folder. + + Moves every entry of ``job_name`` except the optimized structures file + into ``folder``. **Nothing outside ``job_name`` is ever touched**, which + is a correctness requirement and not a style preference: the caller + (``workflow_workers.optim_rank_wrapper``) tars ``folder``, ``rmtree``s it, + and -- under the default ``verbose=False`` -- sends the tarball to trash + or, when that is unavailable (the cluster path), plainly ``os.remove``s + it. Whatever ends up in ``folder`` is therefore *deleted*. + + This function used to additionally sweep ``oeomega_*`` and ``flipper_*`` + out of the **process working directory**, which for an ordinary + ``cd ~/project && auto3d run mols.smi --k 1`` is the user's own directory: + a file named e.g. ``~/project/oeomega_settings.txt`` was moved into the + run's ``verbose`` folder and then destroyed with it, unrecoverably on the + ``os.remove`` path. That loop ran on *every* run, not only OpenEye ones. + The OpenEye logfiles it existed to collect now land inside the chunk + directory instead -- ``isomer_engine.oe_isomer`` runs the OpenEye section + with its working directory set to the directory it owns -- so the loop + below collects them like any other metadata file. + + Each move is guarded individually: a single file that cannot be moved + (permissions, a vanished file) must not abandon the rest of the sweep and + leave a half-populated ``verbose`` folder plus a spurious traceback + behind. Everything here is diagnostic -- the ranked output is excluded and + has already been written by the time this runs. + + Args: + job_name: Path to the job directory containing files to move. + folder: Destination folder for metadata files. + optimized_structures: Path to the final output file (not moved). + + Returns: + None. Moves files to the destination folder. + + Example: + >>> housekeeping("/tmp/job1", "/tmp/job1/verbose", "/tmp/job1/output.sdf") + """ + files = list(Path(job_name).glob("*")) + for file in files: + if str(file) == optimized_structures: + continue + try: + shutil.move(str(file), folder) + except OSError: + logger.warning( + "Could not move %s into %s; leaving it where it is.", file, folder + ) diff --git a/src/Auto3D/models/adapter.py b/src/Auto3D/models/adapter.py index 755e4bda..a6e69942 100644 --- a/src/Auto3D/models/adapter.py +++ b/src/Auto3D/models/adapter.py @@ -369,7 +369,15 @@ class ANI2xtAdapter(BaseModelAdapter): ANI2xt is a retrained version of ANI with improved performance. Uses indexed species (H=0, C=1, N=2, O=3, F=4, S=5, Cl=6). - This model benefits significantly from torch.compile() optimization. + ``compile_model=True`` compiles ``ANI2xt.forward``. Until this change it compiled + *nothing*: ``forward``'s per-element loop contained a data-dependent branch + (``if mask.any():``), and a graph break inside a loop gives Dynamo nowhere to + place a resume point, so it skipped the frame -- measured as **zero** + compiled subgraphs. The loop is now free of data-dependent ops and compiles + to one subgraph (``tests/test_ani2xt_atom_energies.py``). Whether that is a + wall-clock win, and by how much, is a GPU measurement this repository does + not make; see ``benchmarks/bench_optimization_perf.py``. No speedup figure + is claimed here because none has been measured. """ def __init__(self, device: torch.device, compile_model: bool = False) -> None: @@ -384,9 +392,28 @@ def __init__(self, device: torch.device, compile_model: bool = False) -> None: # module scope creates models -> batch_opt -> models, and because # Auto3D/__init__.py eagerly imports Auto3D.batch_opt.ANI2xt_no_rep that # becomes an import cycle at package-import time. - from Auto3D.batch_opt.ANI2xt_no_rep import ANI2xt + from Auto3D.batch_opt.ANI2xt_no_rep import ( + ANI2xt, + element_indices, + self_atomic_energies, + ) model = ANI2xt(device) + num_elements = len(model.networks) + energy_shifts = model.energy_shifts super().__init__(model, device, coord_pad=0.0, species_pad=-1, compile_model=compile_model) + # Precompute-and-pass plumbing for ANI2xt.forward. Both helpers are pure + # functions of `species`, and both have to be called from *outside* + # ANI2xt.forward for it to be compilable at all: element_indices has a + # data-dependent output shape, and a graph break inside forward's + # per-element loop makes Dynamo skip the whole frame rather than split + # it, which is why compile_model=True used to produce zero subgraphs for + # this model. Bound here rather than imported at module scope because + # models -> batch_opt is the one deliberate back-edge and must stay + # inside a method (see the deferred ANI2xt import above). + self._element_indices = element_indices + self._self_atomic_energies = self_atomic_energies + self._num_elements = num_elements + self._energy_shifts = energy_shifts def to_species(self, atomic_numbers: Sequence[int]) -> list[int]: """Remap atomic numbers to ANI2xt's 0-based network indices. @@ -419,7 +446,39 @@ def energy( raises. Energies come out float64 (see ``ANI2xt.forward``); coords are consumed at whatever dtype they arrive in. """ - return self.model(species, coords) + return self._call_model(species, coords) + + def _call_model(self, species: torch.Tensor, coords: torch.Tensor) -> torch.Tensor: + """Invoke ``ANI2xt.forward`` with the species-only terms precomputed. + + ``element_indices`` collapses seven ``nonzero`` calls -- seven + host-device synchronizations per forward on CUDA -- into a single host + readback of a fixed-size count vector, and ``self_atomic_energies`` + removes a seven-iteration Python loop that recomputed a constant. Doing + both out here rather than inside ``forward`` is also what leaves + ``forward`` free of data-dependent ops, so ``compile_model=True`` has a + frame it can actually compile. + + Not cached across calls: the optimization loop gathers a fresh + ``species`` tensor for the still-active subset on every step, so there is + no object whose identity could key a cache, and a content-keyed cache + would cost the comparison it saves. + + Falls back to the plain two-argument call when the helpers are absent, + which happens whenever ``self.model`` is not a real ``ANI2xt`` -- several + tests bypass ``__init__`` and substitute a toy quadratic model so they + can exercise the *real* ``forward``/``energy`` without loading weights or + importing torchani. The precompute is an optimization, not part of the + model contract, so it degrades rather than breaking. + """ + helper = getattr(self, "_element_indices", None) + if helper is None: + return self.model(species, coords) + elem_index = helper(species, self._num_elements) + self_energies = self._self_atomic_energies( + species, self._energy_shifts, self._num_elements) + return self.model(species, coords, elem_index=elem_index, + self_energies=self_energies) def forward( self, @@ -444,7 +503,7 @@ def forward( Tuple of (energies, forces) in eV units. """ coords = coords.requires_grad_(True) - energy = self.model(species, coords) + energy = self._call_model(species, coords) # create_graph=False (default) avoids building second-order gradient graph grad = torch.autograd.grad([energy.sum()], [coords], create_graph=False)[0] forces = -grad @@ -458,7 +517,10 @@ class ANI2xAdapter(BaseModelAdapter): ANI2x uses periodic table indexing for species. Requires torchani to be installed. - This model benefits significantly from torch.compile() optimization. + ``compile_model=True`` compiles the torchani model. No speedup figure is + claimed: none has been measured for this path, and the ~1.25x these + docstrings used to assert had no measurement behind it. + ``benchmarks/bench_optimization_perf.py`` is what produces one. """ def __init__(self, device: torch.device, compile_model: bool = False) -> None: diff --git a/src/Auto3D/processors.py b/src/Auto3D/processors.py index 8e8aa665..51f33b22 100644 --- a/src/Auto3D/processors.py +++ b/src/Auto3D/processors.py @@ -7,9 +7,9 @@ from __future__ import annotations from Auto3D.config import Auto3DOptions -from Auto3D.isomers import create_tautomer_engine -from Auto3D.utils.file_ops import hash_taut_smi +from Auto3D.isomers.factory import create_tautomer_engine from Auto3D.utils.logging_config import get_logger +from Auto3D.utils.smi_io import hash_taut_smi logger = get_logger(__name__) diff --git a/src/Auto3D/ranking.py b/src/Auto3D/ranking.py index c9fde206..694513f0 100644 --- a/src/Auto3D/ranking.py +++ b/src/Auto3D/ranking.py @@ -7,13 +7,18 @@ from Auto3D.config import SELECTOR_FIELDS, check_selectors_mutually_exclusive from Auto3D.exceptions import ConfigurationError, InputValidationError -from Auto3D.filtering import filter_unique_optimized -from Auto3D.utils.chemistry import check_connectivity, ev2kcalpermol, filter_unique +from Auto3D.filtering import filter_unique, filter_unique_optimized +from Auto3D.utils.connectivity import check_connectivity from Auto3D.utils.convergence import converged_or_unfiltered, has_convergence_flag -from Auto3D.utils.energy import E_TOT_HARTREE_PROP, E_TOT_PROP, e_tot_ev +from Auto3D.utils.energy import ( + E_TOT_HARTREE_PROP, + E_TOT_PROP, + e_tot_ev, + ev2kcalpermol, +) from Auto3D.utils.logging_config import get_logger +from Auto3D.utils.output_guard import check_output_not_input, check_output_overwrite from Auto3D.utils.stereo_check import stereo_preserved -from Auto3D.utils.validation import check_output_not_input, check_output_overwrite logger = get_logger(__name__) @@ -32,7 +37,7 @@ def species_id(name: str) -> str: Stripping on the FIRST underscore is wrong whenever ``species_id`` itself contains an underscore -- notably ``smiles2smi``'s InChIKey-collision - disambiguation (``utils/file_ops.py``), which renames a duplicate input's + disambiguation (``utils/smi_io.py``), which renames a duplicate input's id to ``f"{inchikey}_{count}"`` (e.g. ``KEY_2``) specifically so it is not dropped. Stripping the trailing two components with ``rsplit(..., 2)`` instead recovers ``species_id`` intact (embedded underscores and all), so diff --git a/src/Auto3D/tautomer.py b/src/Auto3D/tautomer.py index 6c8be92d..438941e1 100644 --- a/src/Auto3D/tautomer.py +++ b/src/Auto3D/tautomer.py @@ -8,8 +8,7 @@ from Auto3D.auto3D import main from Auto3D.config import Auto3DOptions from Auto3D.exceptions import ConfigurationError -from Auto3D.utils.chemistry import hartree2kcalpermol -from Auto3D.utils.energy import e_tot_hartree +from Auto3D.utils.energy import e_tot_hartree, hartree2kcalpermol from Auto3D.utils.logging_config import get_logger logger = get_logger(__name__) @@ -35,7 +34,7 @@ def select_tautomers(sdf: str, k: int | None = None, window: float | None = None ``auto3d tautomers`` additionally gates its ``-o`` with ``check_output_overwrite`` -- so this is a hazard for direct API callers only, the same residual class as - ``Auto3D.utils.file_ops.smiles2smi`` and ``decode_ids``. See + ``Auto3D.utils.smi_io.smiles2smi`` and ``Auto3D.id_mapping.decode_ids``. See ``docs/superpowers/follow-ups-after-4.0.0-remediation.md``. Note: diff --git a/src/Auto3D/utils/__init__.py b/src/Auto3D/utils/__init__.py index 3f9b21fe..3adcbb90 100644 --- a/src/Auto3D/utils/__init__.py +++ b/src/Auto3D/utils/__init__.py @@ -2,9 +2,9 @@ This package is a **namespace, not a barrel**: it re-exports nothing, and ``__init__.py`` deliberately contains no code at all. Import each name from the -module that defines it -- ``from Auto3D.utils.chemistry import hartree2ev``, not +module that defines it -- ``from Auto3D.utils.energy import hartree2ev``, not ``from Auto3D.utils import hartree2ev``. The barrel that used to live here -listed 41 names drawn from five of the eight modules below, so the three it +listed 41 names drawn from five of the then-eight modules, so the three it omitted (``energy``, ``convergence``, ``stereo_check``) had no way in, and the same function was reached by two different paths in sibling modules. See ``docs/source/api.rst`` for the rule and the CHANGELOG for the full old-to-new @@ -15,19 +15,36 @@ What each module owns: -``chemistry`` - Energy-unit conversions and their constants, molecular properties (charge, - connectivity), RMSD, clash relief, conformer-count heuristics. +``atomic_io`` + The one way this package replaces a file in place: stage a sibling + temporary, copy the target's mode onto it, ``os.replace``. Used wherever a + file is rewritten, so no call site invents its own temp naming again. +``connectivity`` + Bond-graph comparison -- whether optimization preserved connectivity. ``convergence`` The single owner of the ``Converged`` SDF property -- reading it, writing it, and deciding what counts as converged. ``energy`` - The single owner of the ``E_tot`` SDF properties, in both eV and hartree. -``file_ops`` - File I/O: SMILES/SDF reading and writing, SDF chunking, ID encode/decode, - housekeeping, output reordering. + The single owner of the ``E_tot`` SDF properties, in both eV and hartree, + and of the energy-unit constants. +``geometry`` + RMSD and geometric comparison between conformers. ``logging_config`` Logging setup and the logger factory every module calls. +``molprops`` + Molecular properties read off an RDKit mol: charge, and the + conformer-count heuristic. +``output_guard`` + Output-path gates (overwrite, and output-is-not-input). Split out of + ``validation`` so a ``.smi`` writer does not pull in torch and the whole + model tree through it. +``reconciliation`` + Comparing what went in against what came out, so a molecule lost in the + pipeline is visible rather than silent. +``sdf_io`` + SDF reading, writing, chunking and reordering. +``smi_io`` + ``.smi`` reading and writing, and the ID-hashing that goes with it. ``stereo_check`` Species keys and the stereochemistry-preserved check applied after optimization. diff --git a/src/Auto3D/utils/atomic_io.py b/src/Auto3D/utils/atomic_io.py new file mode 100644 index 00000000..4d851c42 --- /dev/null +++ b/src/Auto3D/utils/atomic_io.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python +"""Rewriting a file in place without ever being able to destroy it. + +Three functions in Auto3D read a file and then rewrite it: ``reorder_sdf``, +``ASE.geometry._annotate_and_rewrite`` and +``utils.stereochemistry.amend_configuration_w``. Opening the target directly +truncates it, so a failure partway through the rewrite leaves a half-written +file and no copy of what was there before -- the crash half of audit C14. + +All three grew their own staging code, and the copies diverged. Two used +``mkstemp`` plus a ``chmod`` from the target; ``reorder_sdf`` used a predictable +``.reorder.tmp`` and no ``chmod`` at all, so a 0600 SDF came back at +whatever the process umask allows -- a permission *loosening* on the path an +ordinary ``auto3d run`` takes. This module is the single implementation, so +there is one place for that behavior to be correct. + +Releasing whatever handle the caller opened on the temp path stays the caller's +duty, and matters on Windows: ``os.replace`` refuses a destination another open +handle holds (``PermissionError``/``WinError 5``), and an RDKit +``SDMolSupplier`` on the target is exactly such a handle. +""" +from __future__ import annotations + +import os +import stat +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager + +__all__ = ["atomic_write_path"] + + +@contextmanager +def atomic_write_path( + target: str | os.PathLike[str], *, suffix: str | None = None +) -> Iterator[str]: + """Yield a temp path to write, then atomically move it onto ``target``. + + A **path**, not an open handle, because the writers using this need one: + ``Chem.SDWriter`` takes a filename and opens the file itself. + + On clean exit the temp file replaces ``target`` with ``os.replace``, which + is atomic on POSIX and on Windows -- so ``target`` is only ever the old + complete file or the new complete file, never a partial one. On **any** + exception (``BaseException``, so a ``KeyboardInterrupt`` mid-write counts) + the temp file is removed and the exception propagates, leaving ``target`` + exactly as it was. + + The temp file is created as a **sibling** of ``target``: ``os.replace`` + raises ``OSError: [Errno 18] EXDEV`` across filesystems, and a separate + ``/tmp`` tmpfs is the common case. The parent directory is resolved with + ``realpath``, not ``abspath``: ``abspath`` collapses ``..`` lexically, so a + target like ``/scratch/link/../out.sdf`` (where ``link`` points at another + mount) would stage the temp file in ``/scratch`` while the replace + destination really lives elsewhere, and the ``EXDEV`` would surface only + after the work was finished. Only the PARENT is resolved -- ``os.replace`` + acts on the final path component itself, so following a symlinked + ``target`` would pick the wrong directory. + + The temp file inherits ``target``'s permission bits, set before anything is + written. ``tempfile.mkstemp`` creates 0600, and ``os.replace`` carries the + SOURCE file's mode to the destination, so without this every rewrite would + tighten a 0644 output to 0600 -- and a hand-rolled ``open()`` would loosen a + 0600 target to the umask instead. Setting the mode up front also preserves a + read-only (0444) target's protection, which ``rename(2)`` would otherwise + bypass. Copying the mode is best effort: a target that does not exist yet, + or whose mode cannot be read, is not a reason to abandon a completed + computation. + + Args: + target: The file to replace. Need not exist yet. + suffix: Suffix for the temp file name. Useful only for readability of + a leftover file in a crash dump; the name is otherwise opaque and + unpredictable, which a caller-derived name (``.reorder.tmp``) + is not. + + Yields: + The path of the temp file to write. It exists and is empty. + """ + target_path = os.fspath(target) + directory = os.path.realpath(os.path.dirname(os.path.abspath(target_path))) + fd, tmp_path = tempfile.mkstemp(suffix=suffix, dir=directory) + os.close(fd) + try: + os.chmod(tmp_path, stat.S_IMODE(os.stat(target_path).st_mode)) + except OSError: + # Best effort: `target` may not exist yet (a first write), or its mode + # may be unreadable. Neither is a reason to refuse the rewrite. + pass + + try: + yield tmp_path + os.replace(tmp_path, target_path) + except BaseException: + # BaseException, not Exception: a KeyboardInterrupt mid-write must not + # leave a stray file beside the user's output. + try: + os.unlink(tmp_path) + except OSError: + pass + raise diff --git a/src/Auto3D/utils/chemistry.py b/src/Auto3D/utils/chemistry.py deleted file mode 100644 index 05e7fd54..00000000 --- a/src/Auto3D/utils/chemistry.py +++ /dev/null @@ -1,580 +0,0 @@ -#!/usr/bin/env python -"""Chemistry-related utility functions for Auto3D. - -This module provides: -- Energy unit conversions (Hartree, eV, kcal/mol) -- Molecular charge calculation -- Geometry utilities (pairwise distances, RMSD) -- Molecular connectivity analysis -- RMSD-based duplicate filtering -""" -from __future__ import annotations - -import logging - -import numpy as np -from rdkit import Chem -from rdkit.Chem import AllChem, rdMolAlign, rdMolDescriptors, rdmolops, rdMolTransforms - -from Auto3D.constants import ( - CONFORMER_MULTIPLIER, - CONFORMER_ROTATABLE_COEFF, - CONFORMER_ROTATABLE_EXP, - DEFAULT_DUPLICATE_ENERGY_TOL, - DEFAULT_RMSD_THRESHOLD, - EV_TO_KCAL_PER_MOL, - HARTREE_TO_EV, - HARTREE_TO_KCAL_PER_MOL, - MAX_CONFORMERS_CAP, - MIN_ATOM_DISTANCE, -) -from Auto3D.utils.convergence import converged_or_unfiltered -from Auto3D.utils.energy import try_e_tot_ev -from Auto3D.utils.stereo_check import ( - species_key, - stereo_descriptors_from_3d, - stereo_preserved, -) - -logger = logging.getLogger("auto3d") - -# Re-export constants for convenience -__all__ = [ - # Constants - "HARTREE_TO_EV", - "HARTREE_TO_KCAL_PER_MOL", - "EV_TO_KCAL_PER_MOL", - # Backward compatibility aliases - "hartree2ev", - "hartree2kcalpermol", - "ev2kcalpermol", - # Functions - "calculate_conformer_count", - "get_mol_charge", - "min_pairwise_distance", - "relieve_clash", - "get_rmsd", - "check_connectivity", - "amend_mol", - "get_mol_connectivity", - "filter_unique", -] - -# Backward compatibility aliases for energy conversion factors -hartree2ev: float = HARTREE_TO_EV -hartree2kcalpermol: float = HARTREE_TO_KCAL_PER_MOL -ev2kcalpermol: float = EV_TO_KCAL_PER_MOL - - -def calculate_conformer_count(mol: Chem.Mol) -> int: - """Calculate the number of conformers to generate for a molecule. - - Uses a formula based on the number of rotatable bonds, with a minimum - of the heavy atom count and a maximum cap. The result is floored at 1 so - a molecule never gets 0 conformers (which would silently drop tiny species - such as ``[H+]`` or a lone atom from the pipeline). - - Formula: min(max(1, num_heavy, 2 * 8.481 * (num_rotatable ** 1.642)), 1000) - Reference: https://doi.org/10.1021/acs.jctc.0c01213 - - Args: - mol: RDKit molecule object (with or without hydrogens). - - Returns: - Number of conformers to generate (always >= 1). - - Example: - >>> from rdkit import Chem - >>> mol = Chem.MolFromSmiles("CCCCCC") # hexane - >>> count = calculate_conformer_count(mol) - >>> 1 <= count <= 1000 - True - """ - num_rotatable = rdMolDescriptors.CalcNumRotatableBonds(mol) - num_heavy = sum(1 for atom in mol.GetAtoms() if atom.GetAtomicNum() > 1) - - formula_count = int( - CONFORMER_MULTIPLIER * CONFORMER_ROTATABLE_COEFF * - (num_rotatable ** CONFORMER_ROTATABLE_EXP) - ) - - # Floor at 1: a heavy-atom-free species (e.g. [H+]) or a single atom must - # still receive at least one conformer instead of being silently dropped. - return min(max(1, num_heavy, formula_count), MAX_CONFORMERS_CAP) - - -def get_mol_charge(mol: Chem.Mol) -> int: - """Get the formal charge of a molecule. - - Args: - mol: RDKit Mol object. - - Returns: - The total formal charge of the molecule. - - Example: - >>> from rdkit import Chem - >>> mol = Chem.MolFromSmiles("[NH4+]") - >>> get_mol_charge(mol) - 1 - """ - return rdmolops.GetFormalCharge(mol) - - -def min_pairwise_distance(points: np.ndarray) -> float: - """Find the minimum pairwise distance among n points in 3D space. - - This function computes all pairwise distances between the provided points - and returns the minimum distance. It uses vectorized NumPy operations - for efficiency. - - Args: - points: A (n, 3) array representing the coordinates of n points - in 3D space. - - Returns: - The minimum pairwise distance among the n points. - - Example: - >>> import numpy as np - >>> points = np.array([[0, 0, 0], [1, 0, 0], [0, 2, 0]]) - >>> min_pairwise_distance(points) - 1.0 - """ - # Ensure input is a NumPy array with float32 type - points = points.astype(np.float32) - n = points.shape[0] - - # Guard for single atom or empty input - if n < 2: - # Single atom: no pairwise distance exists - return float('inf') - - # Expand dimensions of points to enable broadcasting - points_expanded = np.expand_dims(points, axis=1).repeat(n, axis=1) - - # Compute pairwise squared differences - diff_squared = (points_expanded - points_expanded.transpose(1, 0, 2)) ** 2 - - # Sum along the last dimension to get pairwise squared distances - pairwise_squared_distances = np.sum(diff_squared, axis=-1) - - # Find the minimum squared distance from upper triangle - upp_indices = np.triu_indices(n, 1) - upp_values = pairwise_squared_distances[upp_indices] - min_squared_distance = np.min(upp_values) - - # Return the square root of the minimum squared distance - return float(np.sqrt(min_squared_distance)) - - -def relieve_clash( - mol: Chem.Mol, - conf_id: int, - min_distance: float = MIN_ATOM_DISTANCE, -) -> bool: - """Optimize a clashing conformer in place and report whether it is usable. - - A conformer is considered "clashing" when its minimum pairwise interatomic - distance is below ``min_distance``. Such conformers are relaxed with MMFF; - when the molecule lacks full MMFF parameters (elements like B, Se or some - Si valences, where ``MMFFOptimizeMolecule`` returns -1 and does nothing), - the function falls back to UFF so the conformer is not discarded for lack - of a force field. - - The force-field relaxation can itself invert a stereocenter or rotate a - double bond. This runs before the enumerated SDF is written, so the - downstream post-optimization stereochemistry check would otherwise read an - already-changed geometry as its own "before" reference and never notice. - Stereochemistry is therefore checked before and after the relaxation, on - this same molecule object, and a conformer whose configuration changed is - rejected here rather than passed downstream. - - Known limitation: the "before" snapshot is read while the conformer is - still in violation of ``min_distance`` — by definition, since that is the - only way execution reaches this branch. CIP perception on a geometry that - is itself clashing is not a trustworthy baseline, unlike the equivalent - check in ``batch_opt/batchopt.py``, whose "before" reading is always - taken from a valid, non-clashing conformer. This matters only when the - branch is actually reached: across roughly 650 conformers sampled from - Auto3D's real ``EmbedMultipleConfs`` output (glucose, cholesterol, a - tripeptide, macrocycles, a cage compound, and molecules with B/Se/ - hypervalent Si), none ever fell below the clash threshold. Under 196 - artificially forced clashes, this guard rejected 96 conformers, and about - 56% of those rejections had a post-relaxation configuration that actually - matched the molecule's true configuration -- spurious rejections caused - by the unreliable baseline rather than a real inversion. The known - improvement is to compare against the molecule's graph-encoded stereo - tags instead of a 3D read of the clashing geometry, but that needs its - own measurement first: RDKit's graph ``AssignStereochemistry`` and - ``AssignStereochemistryFrom3D`` label pseudoasymmetric centers - differently (``r``/``s`` vs ``R``/``S``), which could introduce a - systematic false positive. - - Args: - mol: RDKit molecule holding the conformer. - conf_id: Index of the conformer to check/optimize. - min_distance: Minimum acceptable interatomic distance (Angstroms). - - Returns: - True if the (possibly optimized) conformer's minimum pairwise distance - is >= ``min_distance`` and its stereochemistry survived unchanged; - False if it still clashes or if the relaxation changed its - configuration. - """ - positions = mol.GetConformer(conf_id).GetPositions() - # Closing the dead band: a conformer exactly at the threshold is kept. - if min_pairwise_distance(positions) >= min_distance: - return True - - # Clashing conformer: try MMFF, fall back to UFF when MMFF is unavailable. - before = stereo_descriptors_from_3d(mol, conf_id=conf_id) - if AllChem.MMFFHasAllMoleculeParams(mol): - AllChem.MMFFOptimizeMolecule(mol, confId=conf_id) - else: - AllChem.UFFOptimizeMolecule(mol, confId=conf_id) - - # Clash relief is a force-field relaxation and can invert a center just as - # the neural network optimization can. It runs before the enumerated SDF is - # written, so the post-optimization check downstream would read an already - # inverted geometry as its reference and never notice. Reject the conformer - # here instead; the embedder simply keeps the ones that survive. - if stereo_descriptors_from_3d(mol, conf_id=conf_id) != before: - logger.warning( - "Discarding a conformer whose stereochemistry changed during clash " - "relief." - ) - return False - - positions = mol.GetConformer(conf_id).GetPositions() - return min_pairwise_distance(positions) >= min_distance - - -def get_rmsd(mol1: Chem.Mol, mol2: Chem.Mol, remove_hs: bool = True) -> float: - """Calculate the RMSD between two molecular conformers. - - Uses RDKit's GetBestRMS function which finds the optimal alignment - between the two molecules before computing RMSD. - - Args: - mol1: First RDKit Mol object with a conformer. - mol2: Second RDKit Mol object with a conformer. - remove_hs: If True (default), remove hydrogens before RMSD calculation. - This speeds up the calculation and focuses on heavy atom positions. - - Returns: - The RMSD value in Angstroms. Returns ``float("inf")`` if alignment - fails (e.g., due to atom mismatch). An incomparable pair is treated as - "distinct" rather than "identical", which is the same convention used - by ``filter_unique``; a downstream ``rmsd < threshold`` check therefore - keeps the structure instead of dropping it as a false duplicate. - - Example: - >>> from rdkit import Chem - >>> from rdkit.Chem import AllChem - >>> mol1 = Chem.MolFromSmiles("CCO") - >>> mol1 = Chem.AddHs(mol1) - >>> AllChem.EmbedMolecule(mol1) - 0 - >>> mol2 = Chem.Mol(mol1) # Copy - >>> get_rmsd(mol1, mol2) - 0.0 - """ - try: - if remove_hs: - mol1_proc = Chem.RemoveHs(mol1) - mol2_proc = Chem.RemoveHs(mol2) - else: - mol1_proc = mol1 - mol2_proc = mol2 - # Temporary bug fix for https://github.com/rdkit/rdkit/issues/6826 - rmsd = rdMolAlign.GetBestRMS(mol1_proc, mol2_proc) - except RuntimeError: - # Incomparable pair: treat as distinct (inf), matching filter_unique. - rmsd = float("inf") - return float(rmsd) - - -def check_connectivity(mol: Chem.Mol) -> bool: - """Check if there is a new bond formed or a bond broken in the molecule. - - This function validates molecular connectivity by comparing actual interatomic - distances against reference bond lengths based on UFF radii. It detects both - broken bonds (distances too large) and formed bonds (distances too small). - - Args: - mol: RDKit molecule object with conformer information. - - Returns: - True if connectivity is valid (no broken/formed bonds), False otherwise. - - Note: - Uses UFF bond radii from Rappe et al. JACS 1992. The radii neglect bond-order - and electronegativity corrections. Bond is considered broken if length > 1.25x - reference, and formed if distance < 1.1x reference. - - Bonds involving elements outside the covalent-radii table (e.g. alkali/ - alkaline-earth counterions or transition-metal coordination bonds, M-L) - are NOT validated -- such pairs are skipped ("no opinion"), so the - dissociation of an M-L bond will not be flagged as invalid connectivity. - """ - # Initialize UFF bond radii (Rappe et al. JACS 1992) - # Units of angstroms - # These radii neglect the bond-order and electronegativity corrections in the - # original paper. Where several values exist for the same atom, the largest - # was used. Consequence: a single bond-order-blind reference length makes the - # broken-bond (1.25x) check lenient and the formed-bond (1.1x) check strict, - # so a stretched aromatic/conjugated bond or a short multiple bond can be - # mis-judged. The molecular graph already carries bond orders (see - # get_mol_connectivity); a bond-order-aware reference would be more accurate - # but is intentionally not used here. - Radii = { - 1: 0.354, - 5: 0.838, - 6: 0.757, - 7: 0.700, - 8: 0.658, - 9: 0.668, - 14: 1.117, - 15: 1.117, - 16: 1.064, - 17: 1.044, - 32: 1.197, - 33: 1.211, - 34: 1.190, - 35: 1.192, - 51: 1.407, - 52: 1.386, - 53: 1.382, - } - - atoms = [atom for atom in mol.GetAtoms()] - n = len(atoms) - for i in range(n): - for j in range(i + 1, n, 1): - atom_i = atoms[i] - atom_i_idx = atom_i.GetIdx() - atomic_num_i = atom_i.GetAtomicNum() - pos_i = mol.GetConformer().GetAtomPosition(atom_i_idx) - - atom_j = atoms[j] - atom_j_idx = atom_j.GetIdx() - atomic_num_j = atom_j.GetAtomicNum() - pos_j = mol.GetConformer().GetAtomPosition(atom_j_idx) - - # Elements outside the UFF radii table (e.g. Na, K, Mg, Fe, Zn in - # salts/metal complexes) have no reference radius. Skip such pairs - # ("no opinion") rather than indexing the dict blindly, which would - # raise KeyError and crash the whole filtering pass. - if atomic_num_i not in Radii or atomic_num_j not in Radii: - continue - - bond = mol.GetBondBetweenAtoms(atom_i_idx, atom_j_idx) - reference_length = Radii[atomic_num_i] + Radii[atomic_num_j] - if bond: - # make sure the bond is not broken - length = rdMolTransforms.GetBondLength(mol.GetConformers()[0], atom_i_idx, atom_j_idx) - if length > reference_length * 1.25: - return False - else: - # make sure the bond is not formed - dist = np.linalg.norm(np.array(pos_i) - np.array(pos_j)) - if dist < reference_length * 1.1: - return False - return True - - -def amend_mol( - mol: Chem.Mol, - sanitize: bool = False, - check_valid: bool = False, -) -> Chem.Mol | None: - """Attempt to fix or validate a molecule. - - This function can optionally sanitize a molecule and check its validity. - If check_valid is True and the molecule has invalid connectivity (broken - or formed bonds), None is returned. - - Args: - mol: RDKit Mol object to amend. - sanitize: If True, sanitize the molecule using RDKit's SanitizeMol. - check_valid: If True, check connectivity and return None if invalid. - - Returns: - The amended molecule, or None if the molecule is invalid and check_valid is True. - - Example: - >>> from rdkit import Chem - >>> mol = Chem.MolFromSmiles("CCO") - >>> amended = amend_mol(mol, sanitize=True) - >>> amended is not None - True - """ - if mol is None: - return None - - try: - if sanitize: - Chem.SanitizeMol(mol) - - if check_valid: - # Check if molecule has valid 3D coordinates - if mol.GetNumConformers() > 0: - if not check_connectivity(mol): - return None - - return mol - except (ValueError, RuntimeError, KeyError) as e: - # ValueError: from RDKit SanitizeMol validation errors - # RuntimeError: from RDKit internal errors during molecule processing - # KeyError: defensive only. check_connectivity no longer raises KeyError - # for unknown elements (it now skips them); retained to swallow any - # stray dict-lookup error from RDKit internals rather than crash the - # amendment of a single molecule. - logger.debug(f"Molecule amendment failed: {type(e).__name__}: {e}") - return None - - -def get_mol_connectivity( - mol: Chem.Mol, - include_bond_order: bool = False, -) -> set[tuple[int, int]] | set[tuple[int, int, float]]: - """Get the bond connectivity of a molecule. - - Returns a set of tuples representing bonds in the molecule. Each tuple - contains the indices of the two bonded atoms, optionally with the bond order. - - Args: - mol: RDKit Mol object. - include_bond_order: If True, include bond order as the third element - of each tuple. - - Returns: - A set of tuples. Each tuple is (atom1_idx, atom2_idx) if include_bond_order - is False, or (atom1_idx, atom2_idx, bond_order) if True. The atom indices - are sorted so that atom1_idx < atom2_idx. - - Example: - >>> from rdkit import Chem - >>> mol = Chem.MolFromSmiles("CC") - >>> get_mol_connectivity(mol) - {(0, 1)} - >>> mol = Chem.MolFromSmiles("C=C") - >>> get_mol_connectivity(mol, include_bond_order=True) - {(0, 1, 2.0)} - """ - connectivity: set = set() - - for bond in mol.GetBonds(): - atom1_idx = bond.GetBeginAtomIdx() - atom2_idx = bond.GetEndAtomIdx() - - # Ensure consistent ordering (smaller index first) - if atom1_idx > atom2_idx: - atom1_idx, atom2_idx = atom2_idx, atom1_idx - - if include_bond_order: - bond_order = bond.GetBondTypeAsDouble() - connectivity.add((atom1_idx, atom2_idx, bond_order)) - else: - connectivity.add((atom1_idx, atom2_idx)) - - return connectivity - - -def filter_unique(mols: list[Chem.Mol], crit: float = DEFAULT_RMSD_THRESHOLD) -> list[Chem.Mol]: - """Remove structures that are very similar and remove unconverged structures. - - This function filters a list of molecules to keep only unique, converged structures. - It first removes unconverged structures and those with invalid connectivity, - then removes similar structures based on RMSD comparison. - - Args: - mols: List of RDKit molecule objects, optionally carrying a 'Converged' - property. A record whose 'Converged' is explicitly false is - dropped; a record without the property is kept (not filtered on - convergence). Records marked 'Stereo_changed' are excluded. - crit: RMSD threshold for considering two structures as identical. - Structures with RMSD below this value are considered duplicates. - Defaults to DEFAULT_RMSD_THRESHOLD (0.3 Angstroms). - - Returns: - List of unique, converged molecules with valid connectivity. - - Example: - >>> from rdkit import Chem - >>> from rdkit.Chem import AllChem - >>> mol = Chem.MolFromSmiles("CCO") - >>> mol = Chem.AddHs(mol) - >>> AllChem.EmbedMolecule(mol, randomSeed=42) - 0 - >>> mol.SetProp("Converged", "true") - >>> filter_unique([mol], crit=0.3) # Returns list with 1 molecule - [...] - """ - # Remove structures that explicitly failed to converge. A record with no - # 'Converged' property is NOT filtered on convergence -- see - # Auto3D.utils.convergence for why absence is not failure. - mols_: list[Chem.Mol] = [] - for mol in mols: - convergence_flag = converged_or_unfiltered(mol) - has_valid_bonds = check_connectivity(mol) - if convergence_flag and has_valid_bonds and stereo_preserved(mol): - mols_.append(mol) - mols = mols_ - - # Remove similar structures. Strip Hs once per molecule (O(n)) instead of on - # both sides of every comparison (O(n^2)); GetBestRMS on no-H forms is - # symmetric so results are unchanged. The ORIGINAL (H-explicit) mols are - # returned; no-H forms are comparison-only. - # - # Heavy-atom RMSD alone collapses conformers that differ only in an O-H / N-H - # rotor orientation. Guard with an energy check: a pair counts as duplicate - # only when the RMSD is below ``crit`` AND the energies agree within - # DEFAULT_DUPLICATE_ENERGY_TOL (eV; 'E_tot' is stored in Hartree and is - # converted on read by Auto3D.utils.energy). Mols without a usable 'E_tot' - # fall back to RMSD-only (energy guard cannot apply). - unique_mols: list[Chem.Mol] = [] - unique_noH: list[Chem.Mol] = [] - unique_energies: list[float | None] = [] - unique_species: list[str] = [] - for mol_i in mols: - mol_i_noH = Chem.RemoveHs(mol_i) - # E_tot is stored in Hartree; DEFAULT_DUPLICATE_ENERGY_TOL is in eV. - e_i: float | None = try_e_tot_ev(mol_i) - species_i = species_key(mol_i) - unique = True - for mol_j_noH, e_j, species_j in zip( - unique_noH, unique_energies, unique_species, strict=True - ): - # Two different compounds are never duplicates of each other, however - # close their geometries. All stereoisomers of one input share a - # ranking group, and two ring diastereomers can sit below the default - # 0.3 A threshold, so without this the RMSD test could delete one of - # them (Auto3D.utils.stereo_check.species_key). Checked before the - # RMSD call it makes unnecessary. - if species_i != species_j: - continue - try: - # temporary bug fix for https://github.com/rdkit/rdkit/issues/6826 - # removing Hs speeds up the calculation - rmsd = rdMolAlign.GetBestRMS(mol_i_noH, mol_j_noH) - except RuntimeError: - # Incomparable pair: treat as distinct (not a duplicate) so the - # conformer is kept. Using 0 would make it look like a perfect - # duplicate and drop a genuinely distinct structure. - rmsd = float("inf") - energy_close = ( - e_i is None - or e_j is None - or abs(e_i - e_j) < DEFAULT_DUPLICATE_ENERGY_TOL - ) - if rmsd < crit and energy_close: - unique = False - break - if unique: - unique_mols.append(mol_i) - unique_noH.append(mol_i_noH) - unique_energies.append(e_i) - unique_species.append(species_i) - return unique_mols diff --git a/src/Auto3D/utils/connectivity.py b/src/Auto3D/utils/connectivity.py new file mode 100644 index 00000000..8e154b3d --- /dev/null +++ b/src/Auto3D/utils/connectivity.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python +"""Which atoms are bonded, and whether a 3D geometry still agrees with that. + +:func:`get_mol_connectivity` reads the bonds off the molecular graph; +:func:`check_connectivity` asks whether a conformer's interatomic distances are +consistent with them, which is how a geometry that dissociated or formed a new +bond during optimization is caught. :func:`amend_mol` is the combined +sanitize-and-validate wrapper the filters use. +""" +from __future__ import annotations + +import logging + +import numpy as np +from rdkit import Chem +from rdkit.Chem import rdMolTransforms + +logger = logging.getLogger("auto3d") + +__all__ = ["check_connectivity", "amend_mol", "get_mol_connectivity"] + + +def check_connectivity(mol: Chem.Mol) -> bool: + """Check if there is a new bond formed or a bond broken in the molecule. + + This function validates molecular connectivity by comparing actual interatomic + distances against reference bond lengths based on UFF radii. It detects both + broken bonds (distances too large) and formed bonds (distances too small). + + Args: + mol: RDKit molecule object with conformer information. + + Returns: + True if connectivity is valid (no broken/formed bonds), False otherwise. + + Note: + Uses UFF bond radii from Rappe et al. JACS 1992. The radii neglect bond-order + and electronegativity corrections. Bond is considered broken if length > 1.25x + reference, and formed if distance < 1.1x reference. + + Bonds involving elements outside the covalent-radii table (e.g. alkali/ + alkaline-earth counterions or transition-metal coordination bonds, M-L) + are NOT validated -- such pairs are skipped ("no opinion"), so the + dissociation of an M-L bond will not be flagged as invalid connectivity. + """ + # Initialize UFF bond radii (Rappe et al. JACS 1992) + # Units of angstroms + # These radii neglect the bond-order and electronegativity corrections in the + # original paper. Where several values exist for the same atom, the largest + # was used. Consequence: a single bond-order-blind reference length makes the + # broken-bond (1.25x) check lenient and the formed-bond (1.1x) check strict, + # so a stretched aromatic/conjugated bond or a short multiple bond can be + # mis-judged. The molecular graph already carries bond orders (see + # get_mol_connectivity); a bond-order-aware reference would be more accurate + # but is intentionally not used here. + Radii = { + 1: 0.354, + 5: 0.838, + 6: 0.757, + 7: 0.700, + 8: 0.658, + 9: 0.668, + 14: 1.117, + 15: 1.117, + 16: 1.064, + 17: 1.044, + 32: 1.197, + 33: 1.211, + 34: 1.190, + 35: 1.192, + 51: 1.407, + 52: 1.386, + 53: 1.382, + } + + atoms = [atom for atom in mol.GetAtoms()] + n = len(atoms) + for i in range(n): + for j in range(i + 1, n, 1): + atom_i = atoms[i] + atom_i_idx = atom_i.GetIdx() + atomic_num_i = atom_i.GetAtomicNum() + pos_i = mol.GetConformer().GetAtomPosition(atom_i_idx) + + atom_j = atoms[j] + atom_j_idx = atom_j.GetIdx() + atomic_num_j = atom_j.GetAtomicNum() + pos_j = mol.GetConformer().GetAtomPosition(atom_j_idx) + + # Elements outside the UFF radii table (e.g. Na, K, Mg, Fe, Zn in + # salts/metal complexes) have no reference radius. Skip such pairs + # ("no opinion") rather than indexing the dict blindly, which would + # raise KeyError and crash the whole filtering pass. + if atomic_num_i not in Radii or atomic_num_j not in Radii: + continue + + bond = mol.GetBondBetweenAtoms(atom_i_idx, atom_j_idx) + reference_length = Radii[atomic_num_i] + Radii[atomic_num_j] + if bond: + # make sure the bond is not broken + length = rdMolTransforms.GetBondLength(mol.GetConformers()[0], atom_i_idx, atom_j_idx) + if length > reference_length * 1.25: + return False + else: + # make sure the bond is not formed + dist = np.linalg.norm(np.array(pos_i) - np.array(pos_j)) + if dist < reference_length * 1.1: + return False + return True + + +def amend_mol( + mol: Chem.Mol, + sanitize: bool = False, + check_valid: bool = False, +) -> Chem.Mol | None: + """Attempt to fix or validate a molecule. + + This function can optionally sanitize a molecule and check its validity. + If check_valid is True and the molecule has invalid connectivity (broken + or formed bonds), None is returned. + + Args: + mol: RDKit Mol object to amend. + sanitize: If True, sanitize the molecule using RDKit's SanitizeMol. + check_valid: If True, check connectivity and return None if invalid. + + Returns: + The amended molecule, or None if the molecule is invalid and check_valid is True. + + Example: + >>> from rdkit import Chem + >>> mol = Chem.MolFromSmiles("CCO") + >>> amended = amend_mol(mol, sanitize=True) + >>> amended is not None + True + """ + if mol is None: + return None + + try: + if sanitize: + Chem.SanitizeMol(mol) + + if check_valid: + # Check if molecule has valid 3D coordinates + if mol.GetNumConformers() > 0: + if not check_connectivity(mol): + return None + + return mol + except (ValueError, RuntimeError, KeyError) as e: + # ValueError: from RDKit SanitizeMol validation errors + # RuntimeError: from RDKit internal errors during molecule processing + # KeyError: defensive only. check_connectivity no longer raises KeyError + # for unknown elements (it now skips them); retained to swallow any + # stray dict-lookup error from RDKit internals rather than crash the + # amendment of a single molecule. + logger.debug(f"Molecule amendment failed: {type(e).__name__}: {e}") + return None + + +def get_mol_connectivity( + mol: Chem.Mol, + include_bond_order: bool = False, +) -> set[tuple[int, int]] | set[tuple[int, int, float]]: + """Get the bond connectivity of a molecule. + + Returns a set of tuples representing bonds in the molecule. Each tuple + contains the indices of the two bonded atoms, optionally with the bond order. + + Args: + mol: RDKit Mol object. + include_bond_order: If True, include bond order as the third element + of each tuple. + + Returns: + A set of tuples. Each tuple is (atom1_idx, atom2_idx) if include_bond_order + is False, or (atom1_idx, atom2_idx, bond_order) if True. The atom indices + are sorted so that atom1_idx < atom2_idx. + + Example: + >>> from rdkit import Chem + >>> mol = Chem.MolFromSmiles("CC") + >>> get_mol_connectivity(mol) + {(0, 1)} + >>> mol = Chem.MolFromSmiles("C=C") + >>> get_mol_connectivity(mol, include_bond_order=True) + {(0, 1, 2.0)} + """ + connectivity: set = set() + + for bond in mol.GetBonds(): + atom1_idx = bond.GetBeginAtomIdx() + atom2_idx = bond.GetEndAtomIdx() + + # Ensure consistent ordering (smaller index first) + if atom1_idx > atom2_idx: + atom1_idx, atom2_idx = atom2_idx, atom1_idx + + if include_bond_order: + bond_order = bond.GetBondTypeAsDouble() + connectivity.add((atom1_idx, atom2_idx, bond_order)) + else: + connectivity.add((atom1_idx, atom2_idx)) + + return connectivity diff --git a/src/Auto3D/utils/convergence.py b/src/Auto3D/utils/convergence.py index dbf79652..68b2505b 100644 --- a/src/Auto3D/utils/convergence.py +++ b/src/Auto3D/utils/convergence.py @@ -5,7 +5,7 @@ it optimizes -- ``"True"`` or ``"False"`` -- and read by the three filters that decide which conformers survive (:class:`Auto3D.ranking.ConformerRanker`, :func:`Auto3D.filtering.filter_unique_optimized`, -:func:`Auto3D.utils.chemistry.filter_unique`). +:func:`Auto3D.filtering.filter_unique`). All three used to read it as:: diff --git a/src/Auto3D/utils/energy.py b/src/Auto3D/utils/energy.py index f0d10b61..76fc2787 100644 --- a/src/Auto3D/utils/energy.py +++ b/src/Auto3D/utils/energy.py @@ -10,7 +10,8 @@ ``E_tot`` in eV and ``ASE/geometry.opt_geometry`` converted the same tag to Hartree afterwards, so the identical property name carried two units depending on which entry point produced the file -- and the five in-package consumers -(``ranking``, ``filtering``, ``utils.chemistry.filter_unique``) all hard-coded +(``ranking``, ``filtering.filter_unique_optimized``, +``filtering.filter_unique``) all hard-coded eV. Feeding an ``opt_geometry`` output straight to ``ConformerRanker(window=2.0)`` therefore opened a window 27.2x too wide, kept 3 conformers where 2 belong, reported ``E_rel`` 0.037 kcal/mol where the truth @@ -29,7 +30,11 @@ from rdkit import Chem -from Auto3D.constants import HARTREE_TO_EV +from Auto3D.constants import ( + EV_TO_KCAL_PER_MOL, + HARTREE_TO_EV, + HARTREE_TO_KCAL_PER_MOL, +) __all__ = [ "E_TOT_PROP", @@ -38,6 +43,13 @@ "e_tot_hartree", "e_tot_ev", "try_e_tot_ev", + # Conversion factors, and their lowercase legacy spellings + "HARTREE_TO_EV", + "HARTREE_TO_KCAL_PER_MOL", + "EV_TO_KCAL_PER_MOL", + "hartree2ev", + "hartree2kcalpermol", + "ev2kcalpermol", ] #: Unlabeled property name, kept for backward compatibility. Hartree. @@ -45,6 +57,14 @@ #: Unit-labeled sibling carrying the identical value. E_TOT_HARTREE_PROP = "E_tot(Hartree)" +# Legacy lowercase spellings of the three conversion factors in +# ``Auto3D.constants``. They are the names Auto3D 2.x used and several call +# sites still read, so they stay -- here rather than in a "chemistry" grab bag, +# since a unit conversion factor belongs with the module that owns the unit. +hartree2ev: float = HARTREE_TO_EV +hartree2kcalpermol: float = HARTREE_TO_KCAL_PER_MOL +ev2kcalpermol: float = EV_TO_KCAL_PER_MOL + def set_e_tot_from_ev(mol: Chem.Mol, energy_ev: float, *, labeled: bool = True) -> None: """Write ``E_tot`` (Hartree) from an energy the model produced in eV. diff --git a/src/Auto3D/utils/file_ops.py b/src/Auto3D/utils/file_ops.py deleted file mode 100644 index b1ed1be7..00000000 --- a/src/Auto3D/utils/file_ops.py +++ /dev/null @@ -1,849 +0,0 @@ -#!/usr/bin/env python -"""File operation utilities for Auto3D. - -This module provides functions for file I/O operations, including: -- SMILES file manipulation (hashing IDs, combining files) -- SDF file chunking -- ID encoding/decoding -- Temporary file housekeeping -- SDF reordering -- File type detection -- SMILES encoding for filenames -""" -from __future__ import annotations - -import collections -import os -import shutil -from collections import defaultdict -from pathlib import Path - -from rdkit import Chem -from rdkit.Chem import inchi - -from Auto3D.exceptions import ConfigurationError, InputValidationError -from Auto3D.utils.logging_config import get_logger - -logger = get_logger(__name__) - -#: Stand-in ID for an input record Auto3D could not parse, used by the -#: input-vs-output reconciliation. Such a record has no ``_Name`` to report, so it -#: is identified by its position in the source file. The angle brackets make it -#: unmistakable as a placeholder rather than a molecule name a user chose -- an ID -#: read from a file can be anything, including ``record 3``, but not with these. -UNPARSEABLE_RECORD_ID = "" - - -def iter_smi_records(path, *, on_malformed="skip"): - """Yield (line_no, smiles, mol_id) for each non-blank, non-comment line of - a .smi file. - - A line is 'SMILES ID [extra columns ignored]'. Blank/whitespace-only lines - are skipped, as are lines whose first non-whitespace character is '#' - (comments) -- matching cli.commands.validate.validate_smiles_file, so - `auto3d validate` and every consumer of this function (encode_ids and so - the whole run pipeline, plus the isomer/tautomer engines and the - input/output reconciliation helpers) agree on what a comment line is - (M25). A real SMILES token can never start with '#' (it is a bond symbol - between two atoms, never a leading character), so this cannot misclassify - a legitimate SMILES as a comment. on_malformed controls lines with fewer - than 2 whitespace tokens: - - "skip": log a warning and skip the line (lenient; default) - - "raise": raise InputValidationError naming the line - - Args: - path: Path to the input .smi file. - on_malformed: How to handle lines with fewer than 2 tokens - ("skip" or "raise"). - - Yields: - Tuples of (line_no, smiles, mol_id) where line_no is 1-based. Any extra - whitespace-separated columns beyond the ID are intentionally ignored. - - Raises: - InputValidationError: If on_malformed == "raise" and a non-blank, - non-comment line has fewer than 2 whitespace tokens. - ValueError: If on_malformed is not "skip" or "raise". - """ - if on_malformed not in ("skip", "raise"): - raise ValueError( - f"on_malformed must be 'skip' or 'raise', got: {on_malformed!r}" - ) - with open(path) as f: - data = f.readlines() - for line_no, line in enumerate(data, start=1): - stripped = line.strip() - if not stripped: - continue - if stripped.startswith("#"): - continue - parts = stripped.split() - if len(parts) < 2: - if on_malformed == "raise": - raise InputValidationError( - f"Line {line_no} is missing a molecule ID " - f"(expected 'SMILES ID'): {line.strip()!r}" - ) - logger.warning( - f"Skipping molecule at line {line_no}: failed to parse " - f"(need 'SMILES ID', got: {line.strip()!r})" - ) - continue - # Lenient parsing: ignore any extra whitespace-separated columns. - yield line_no, parts[0], parts[1] - - -def smiles2smi(smiles: list[str], path: str) -> str: - """Convert a list of SMILES strings to a .smi file with InChIKey IDs. - - Each SMILES string is converted to a molecule, and its InChIKey is computed - to serve as a unique identifier. The output file contains one molecule per - line in the format: "SMILES InChIKey". - - Args: - smiles: List of SMILES strings to convert. - path: Output file path for the .smi file. - - Returns: - The output file path. - - Example: - >>> smiles2smi(["CCO", "CCC"], "molecules.smi") - 'molecules.smi' - # File content: - # CCO LFQSCWFLJHTTHZ-UHFFFAOYSA-N - # CCC ATUOYWHBWRKTHZ-UHFFFAOYSA-N - """ - lines = [] - seen_ids: dict[str, int] = {} - for idx, smi in enumerate(smiles): - mol = Chem.MolFromSmiles(smi) - if mol is None: - raise InputValidationError( - f"Invalid SMILES at index {idx}: {smi!r} could not be parsed " - "by RDKit." - ) - inchikey = inchi.MolToInchiKey(mol) - # Distinct inputs can share a standard InChIKey (e.g. tautomers the - # standard InChIKey conflates, or the same molecule written two ways). - # The InChIKey is used as the molecule's unique ID downstream, and - # reorder_sdf collapses duplicate IDs -- so a colliding input would be - # silently dropped. Disambiguate by suffixing repeats (_2, _3, ...) so - # every input keeps its own conformers. The suffix stays a single - # whitespace-delimited token and round-trips through enumeration. - count = seen_ids.get(inchikey, 0) + 1 - seen_ids[inchikey] = count - mol_id = inchikey if count == 1 else f"{inchikey}_{count}" - if count > 1: - logger.info( - "Input SMILES %r shares InChIKey %s with an earlier input; " - "assigning disambiguated id %s so it is not dropped.", - smi, - inchikey, - mol_id, - ) - lines.append(f"{smi} {mol_id}\n") - - with open(path, "w+") as f: - for line in lines: - f.write(line) - - return path - - -def guess_file_type(filename: str) -> str: - """Return the file extension for a given filename. - - Determines the file type based on the extension of the provided filename. - The extension is returned without the leading dot. - - Args: - filename: Path or filename to analyze. - - Returns: - The file extension without the leading dot (e.g., 'smi', 'sdf', 'xyz'). - - Example: - >>> guess_file_type("molecules.sdf") - 'sdf' - >>> guess_file_type("/path/to/input.smi") - 'smi' - >>> guess_file_type("file.mol2") - 'mol2' - """ - return Path(filename).suffix[1:] - - -def hash_enumerated_smi_IDs(smi: str, out: str) -> None: - """Write all SMILES with hashed IDs into a new file. - - Reads a SMILES file, sorts entries by ID, handles duplicate IDs by appending - '_0' suffix, and writes the result to the output file. - - Args: - smi: Path to the input .smi file. - out: Path for the output .smi file with sorted/hashed IDs. - - Returns: - None. Writes the result to the output file. - - Example: - >>> hash_enumerated_smi_IDs("input.smi", "output.smi") - """ - dict0: dict[str, str] = {} - for _line_no, smiles, id in iter_smi_records(smi, on_malformed="skip"): - while id in dict0: - id += "_0" - dict0[id] = smiles - - dict0 = collections.OrderedDict(sorted(dict0.items())) - - with open(out, "w+") as f: - for id, smiles in dict0.items(): - molecule = smiles.strip() + " " + id.strip() + "\n" - f.write(molecule) - - -def hash_taut_smi(smi: str, out: str) -> None: - """Write all SMILES with hashed IDs for tautomers. - - Reads a SMILES file and appends '@tautN' suffix to IDs where N is - an incrementing counter, ensuring unique tautomer identifiers. - - Args: - smi: Path to the input .smi file. - out: Path for the output .smi file with tautomer IDs. - - Returns: - None. Writes the result to the output file. - - Example: - >>> hash_taut_smi("input.smi", "tautomers.smi") - """ - dict0: dict[str, str] = {} - for _line_no, smiles, id in iter_smi_records(smi, on_malformed="skip"): - c = 1 - id_ = id - while ("taut" not in id_) or (id_ in dict0): - id_ = id + f"@taut{c}" - c += 1 - dict0[id_] = smiles - - dict0 = collections.OrderedDict(sorted(dict0.items())) - - with open(out, "w+") as f: - for id, smiles in dict0.items(): - molecule = smiles.strip() + " " + id.strip() + "\n" - f.write(molecule) - - -def housekeeping(job_name: str, folder: str, optimized_structures: str) -> None: - """Move this job directory's metadata files into a folder. - - Moves every entry of ``job_name`` except the optimized structures file - into ``folder``. **Nothing outside ``job_name`` is ever touched**, which - is a correctness requirement and not a style preference: the caller - (``workflow_workers.optim_rank_wrapper``) tars ``folder``, ``rmtree``s it, - and -- under the default ``verbose=False`` -- sends the tarball to trash - or, when that is unavailable (the cluster path), plainly ``os.remove``s - it. Whatever ends up in ``folder`` is therefore *deleted*. - - This function used to additionally sweep ``oeomega_*`` and ``flipper_*`` - out of the **process working directory**, which for an ordinary - ``cd ~/project && auto3d run mols.smi --k 1`` is the user's own directory: - a file named e.g. ``~/project/oeomega_settings.txt`` was moved into the - run's ``verbose`` folder and then destroyed with it, unrecoverably on the - ``os.remove`` path. That loop ran on *every* run, not only OpenEye ones. - The OpenEye logfiles it existed to collect now land inside the chunk - directory instead -- ``isomer_engine.oe_isomer`` runs the OpenEye section - with its working directory set to the directory it owns -- so the loop - below collects them like any other metadata file. - - Each move is guarded individually: a single file that cannot be moved - (permissions, a vanished file) must not abandon the rest of the sweep and - leave a half-populated ``verbose`` folder plus a spurious traceback - behind. Everything here is diagnostic -- the ranked output is excluded and - has already been written by the time this runs. - - Args: - job_name: Path to the job directory containing files to move. - folder: Destination folder for metadata files. - optimized_structures: Path to the final output file (not moved). - - Returns: - None. Moves files to the destination folder. - - Example: - >>> housekeeping("/tmp/job1", "/tmp/job1/verbose", "/tmp/job1/output.sdf") - """ - files = list(Path(job_name).glob("*")) - for file in files: - if str(file) == optimized_structures: - continue - try: - shutil.move(str(file), folder) - except OSError: - logger.warning( - "Could not move %s into %s; leaving it where it is.", file, folder - ) - - -def create_chunk_meta_names(path: str, dir: str) -> dict[str, str]: - """Create output file names based on chunk input path and directory. - - Generates a dictionary of standardized file paths for all intermediate - and output files used in the Auto3D workflow. - - Args: - path: Chunk input .smi file path. - dir: Chunk job folder path. - - Returns: - Dictionary mapping meta names to file paths with the following keys: - - output: Final 3D structure output file - - optimized_og: Original optimized structures - - output_taut: Tautomer SMILES output - - smiles_enumerated: Enumerated SMILES file - - smiles_reduced: Reduced enumerated SMILES file - - smiles_hashed: Hashed enumerated SMILES file - - enumerated_sdf: Enumerated SDF file - - sorted_sdf: Sorted SDF file - - housekeeping_folder: Verbose output folder - - path: Original input path - - dir: Job directory - - Example: - >>> meta = create_chunk_meta_names("chunk1.smi", "/tmp/job") - >>> meta["output"] - '/tmp/job/chunk1_3d.sdf' - """ - dct: dict[str, str] = {} - dir_path = Path(dir) - stem = Path(path).stem - - output = str(dir_path / f"{stem}_3d.sdf") - optimized_og = str(dir_path / f"{stem}_3d0.sdf") - output_taut = str(dir_path / "smi_taut.smi") - smiles_enumerated = str(dir_path / "smiles_enumerated.smi") - smiles_reduced = str(dir_path / "smiles_enumerated_reduced.smi") - smiles_hashed = str(dir_path / "smiles_enumerated_hashed.smi") - enumerated_sdf = str(dir_path / "smiles_enumerated.sdf") - sorted_sdf = str(dir_path / "enumerated_sorted.sdf") - housekeeping_folder = str(dir_path / "verbose") - - dct["output"] = output - dct["optimized_og"] = optimized_og - dct["output_taut"] = output_taut - dct["smiles_enumerated"] = smiles_enumerated - dct["smiles_reduced"] = smiles_reduced - dct["smiles_hashed"] = smiles_hashed - dct["enumerated_sdf"] = enumerated_sdf - dct["sorted_sdf"] = sorted_sdf - dct["housekeeping_folder"] = housekeeping_folder - dct["path"] = path - dct["dir"] = dir - return dct - - -def combine_smi(smies: list[str], out: str) -> None: - """Combine multiple SMILES files into a single file. - - Reads all input SMILES files, removes duplicates, and writes the - combined unique entries to the output file. - - Args: - smies: List of paths to input .smi files. - out: Path for the combined output .smi file. - - Returns: - None. Writes the combined result to the output file. - - Example: - >>> combine_smi(["file1.smi", "file2.smi"], "combined.smi") - """ - data: list[str] = [] - for smi in smies: - with open(smi) as f: - datai = f.readlines() - data += datai - # Order-preserving dedup: list(set(...)) randomizes line order across runs - # (hash seed), making the combined output non-deterministic. dict.fromkeys - # keeps first-seen order while removing exact duplicates. - data = list(dict.fromkeys(data)) - with open(out, "w+") as f2: - for line in data: - if not line.isspace(): - f2.write(line.strip() + "\n") - - -def SDF2chunks(sdf: str) -> list[list[str]]: - """Split an SDF file into chunks, one per molecule. - - Reads an SDF file and splits it into a list of chunks, where each chunk - contains the lines of a single molecule as they appear in the original file. - - Args: - sdf: Path to the input SDF file. - - Returns: - List of chunks, where each chunk is a list of strings (lines) - representing one molecule including the '$$$$' terminator. - - Example: - >>> chunks = SDF2chunks("molecules.sdf") - >>> len(chunks) # Number of molecules - 10 - >>> chunks[0][-1].strip() # Last line of first molecule - '$$$$' - """ - chunks: list[list[str]] = [] - with open(sdf) as f: - data = f.readlines() - chunk: list[str] = [] - for line in data: - if line.strip() == "$$$$": - chunk.append(line) - chunks.append(chunk) - chunk = [] - else: - chunk.append(line) - # A final record lacking the '$$$$' terminator leaves residual lines in - # `chunk`. Preserve it as the last chunk rather than silently dropping it. - if any(line.strip() for line in chunk): - logger.warning( - "SDF file %s ends without a '$$$$' terminator; " - "keeping the trailing record as a final chunk.", - sdf, - ) - chunks.append(chunk) - return chunks - - -def encode_ids( - path: str, out_dir: str | os.PathLike[str] | None = None -) -> tuple[str, dict[str, int]]: - """Encode molecule IDs to numeric indices. - - For a .smi or .sdf file, replaces all molecule IDs with sequential - integer indices and returns a mapping from original IDs to indices. - - The encoded file is named ``_encoded.``. That name is derived - from the input, so it can collide with a file the user already owns: - ``mols_encoded.smi`` sitting beside ``mols.smi`` is a perfectly ordinary - thing for a user to have, and this function used to overwrite it without - a word (``WorkflowOrchestrator`` then ``unlink()``ed it at the end of the - run, so the file was destroyed twice over). Two things prevent that now: - ``out_dir`` lets the caller redirect the encoded file somewhere it owns - -- ``WorkflowOrchestrator`` passes its freshly created job directory -- - and the collision check below refuses to write over an existing file for - every caller, including ones that take the default location. - - Args: - path: Path to the input .smi or .sdf file. - out_dir: Directory to write the encoded file into. Defaults to the - input file's own directory. - - Returns: - Tuple containing: - - Path to the new file with encoded IDs (adds '_encoded' suffix) - - Dictionary mapping original IDs to their numeric indices - - Raises: - ValueError: If the input file is neither .smi nor .sdf format. - ConfigurationError: If a file already exists at the encoded path. - InputValidationError: If a molecule has a missing/blank ID or a - duplicate ID is encountered. - - Example: - >>> new_path, mapping = encode_ids("molecules.smi") - >>> mapping - {'mol_A': 0, 'mol_B': 1, 'mol_C': 2} - """ - path_obj = Path(path).resolve() - extension = path_obj.suffix[1:] - # Checked up front rather than in a trailing `else`: the collision check - # below must not be the thing that reports an unsupported extension. - if extension not in ("smi", "sdf"): - raise ValueError("The input file should be either smi or sdf") - - directory = Path(out_dir) if out_dir is not None else path_obj.parent - new_path = directory / f"{path_obj.stem}_encoded.{extension}" - if new_path.exists(): - raise ConfigurationError( - f"encode_ids would overwrite the existing file {new_path}. " - "Auto3D writes its encoded copy of the input there; move or " - "rename that file, or pass out_dir to write the encoded copy " - "somewhere else." - ) - - if extension == "smi": - new_data: list[str] = [] - mapping: dict[str, int] = {} - # iter_smi_records raises InputValidationError on a <2-token line - # (on_malformed="raise"). Duplicate-id detection stays here because the - # helper does not dedup. Index by a dense record counter, not the file - # line number: blank/skipped lines would otherwise leave gaps in the - # index space, which is inconsistent with the dense positions the chunk - # manager assumes downstream. The original file line_no is still used in - # the error message so it points at the real offending line. - for i, (line_no, smi, id) in enumerate( - iter_smi_records(path, on_malformed="raise") - ): - if id in mapping: - raise InputValidationError( - f"Duplicate molecule ID {id!r} on line {line_no}. " - "IDs must be unique." - ) - mapping[id] = i - new_data.append(f"{smi} {i}\n") - with open(new_path, "w") as f: - for line in new_data: - f.write(line) - return str(new_path), mapping - - else: # "sdf" -- the only remaining possibility, checked above - suppl = Chem.SDMolSupplier(path, removeHs=False) - mapping = {} - with Chem.SDWriter(str(new_path)) as w: - for i, mol in enumerate(suppl): - if mol is None: - logger.warning(f"Skipping molecule at index {i}: failed to parse") - continue - id = mol.GetProp("_Name").strip() - if not id: - raise InputValidationError( - f"Molecule at index {i} has a missing or blank name." - ) - if id in mapping: - raise InputValidationError( - f"Duplicate molecule name {id!r} at index {i}. " - "Names must be unique." - ) - mapping[id] = i - mol.SetProp("_Name", str(i)) - w.write(mol) - return str(new_path), mapping - - -def decode_ids(path: str, mapping: dict[str, int]) -> str: - """Decode numeric IDs back to original molecule IDs. - - For an SDF file with numeric IDs, restores the original IDs using - the provided mapping dictionary. - - Args: - path: Path to the input SDF file with encoded (numeric) IDs. - mapping: Dictionary mapping original IDs to their numeric indices - (as returned by encode_ids). - - Returns: - Path to the new SDF file with decoded IDs (adds '_out' suffix). - - Example: - >>> mapping = {'mol_A': 0, 'mol_B': 1} - >>> output_path = decode_ids("encoded_3d.sdf", mapping) - """ - # Invert the mapping: index -> original_id - inverse_mapping = {v: k for k, v in mapping.items()} - path_obj = Path(path).resolve() - extension = path_obj.suffix[1:] - # Reconstruct base name: remove last two underscore-separated parts - stem_parts = path_obj.stem.split("_")[:-2] - new_stem = "_".join(stem_parts) + "_out" - new_path = path_obj.parent / f"{new_stem}.{extension}" - - suppl = Chem.SDMolSupplier(path, removeHs=False) - with Chem.SDWriter(str(new_path)) as w: - for i, mol in enumerate(suppl): - if mol is None: - logger.warning("Skipping molecule at index %d: failed to parse", i) - continue - name = mol.GetProp("_Name").strip() - if "@taut" in name: - components = name.split("@taut") - new_name = ( - inverse_mapping[int(components[0])] + "@taut" + "".join(components[1:]) - ) - else: - new_name = inverse_mapping[int(name)] - mol.SetProp("_Name", new_name) - - id = "_".join(mol.GetProp("ID").strip().split("_")[1:]) - new_id = new_name + "_" + id - mol.SetProp("ID", new_id) - - w.write(mol) - return str(new_path) - - -def reorder_sdf(sdf: str, source: str) -> list[Chem.Mol]: - """Reorder conformers in an SDF file to match the input source file order. - - Reads the order of molecule IDs from the source file and rewrites the SDF - file with conformers ordered to match. This ensures consistent output - ordering regardless of processing order. - - Args: - sdf: Path to the SDF file to reorder (will be overwritten). - source: Path to the source .smi or .sdf file defining the desired order. - - Returns: - List of RDKit Mol objects in the reordered sequence. - - Note: - - For tautomer conformers (containing '@taut' in ID), the base ID - is extracted for ordering purposes. - - If the source format is unsupported, prints a message and returns None. - - Molecules whose id is not present in ``source`` are appended at the - end (not dropped), so no data is lost. - - Duplicate source ids are de-duplicated: each id's molecules are - written once, so the returned list may be shorter than the input if - source ids repeat. - - Example: - >>> ordered_mols = reorder_sdf("output_3d.sdf", "input.smi") - >>> len(ordered_mols) - 10 - """ - # convert smi/sdf to a list of ids with correct order - ids: list[str] = [] - format = guess_file_type(source) - if format == "smi": - for _line_no, _smiles, mol_id in iter_smi_records(source, on_malformed="skip"): - ids.append(mol_id) - elif format == "sdf": - supp = Chem.SDMolSupplier(source, removeHs=False) - for i, mol in enumerate(supp): - if mol is None: - logger.warning("Skipping molecule at index %d: failed to parse", i) - continue - ids.append(mol.GetProp("_Name")) - else: - logger.warning("Unsupported file format: %s" % format) - return None # type: ignore - - # convert sdf to a Dict[id, List[mols]], preserving discovery order so any - # molecule whose id is not in `source` can still be appended (no data loss). - id_mols: dict[str, list[Chem.Mol]] = defaultdict(lambda: []) - discovery_order: list[str] = [] - supp = Chem.SDMolSupplier(sdf, removeHs=False) - for i, mol in enumerate(supp): - if mol is None: - logger.warning("Skipping molecule at index %d: failed to parse", i) - continue - id = mol.GetProp("_Name") - if "@taut" in id: - id = id.split("@taut")[0] - if id not in id_mols: - discovery_order.append(id) - id_mols[id].append(mol) - - # Release the RDKit supplier's file handle before overwriting `sdf`. - # On Windows an open handle makes the later os.replace() fail with - # "Access is denied" (WinError 5); on POSIX the replace would succeed. - del supp - - # Order: ids present in `source` first (in source order), then any - # unmatched molecules appended in their original order so nothing is lost. - source_id_set = set(ids) - ordered_ids = list(ids) - for id in discovery_order: - if id not in source_id_set: - logger.warning( - "Molecule id %r in %s is not present in source %s; " - "appending it at the end to avoid data loss.", - id, - sdf, - source, - ) - ordered_ids.append(id) - - # write the mols in the correct order to a temp file, then atomically - # replace the original only on success (crash-safe in-place overwrite). - sdf_path = Path(sdf) - tmp_path = sdf_path.with_name(sdf_path.name + ".reorder.tmp") - ordered_mols: list[Chem.Mol] = [] - written_ids: set[str] = set() - try: - with Chem.SDWriter(str(tmp_path)) as f: - for id in ordered_ids: - if id in written_ids: - continue - written_ids.add(id) - mols = id_mols[id] - if len(mols) >= 1: - ordered_mols.extend(mols) - for mol in mols: - f.write(mol) - os.replace(str(tmp_path), str(sdf)) - except BaseException: - # Never leave a half-written temp file or destroy the original input. - try: - tmp_path.unlink() - except OSError: - pass - raise - return ordered_mols - - -def count_sdf(sdf: str) -> int: - """Count the number of molecules in an SDF file. - - Args: - sdf: Path to the SDF file. - - Returns: - Number of molecules in the file. - - Example: - >>> count_sdf("molecules.sdf") - 10 - """ - mols = Chem.SDMolSupplier(sdf) - return len([mol for mol in mols if mol is not None]) - - -def find_smiles_not_in_sdf(smi: str, sdf: str) -> list[tuple[str, str]]: - """Find SMILES that failed to generate 3D conformers. - - Compares a SMILES input file against an SDF output file to identify - molecules that did not successfully generate 3D structures. - - Args: - smi: Path to input SMILES file. - sdf: Path to output SDF file. - - Returns: - List of (id, smiles) tuples for molecules not in SDF. - - Example: - >>> bad = find_smiles_not_in_sdf("input.smi", "output.sdf") - >>> for mol_id, smiles in bad: - ... print(f"Failed: {mol_id}") - """ - # Find all SMILES ids - smi_names: list[tuple[str, str]] = [] - for _line_no, smiles_str, mol_id in iter_smi_records(smi, on_malformed="skip"): - smi_names.append((smiles_str, mol_id)) - - # Get all molecule names from SDF - sdf_data: list[str] = [] - mols = Chem.SDMolSupplier(sdf) - for i, mol in enumerate(mols): - if mol is None: - logger.warning("Skipping molecule at index %d: failed to parse", i) - continue - name = mol.GetProp("_Name") - # decode_ids keeps a "@tautN" suffix on tautomer-enumerated conformers - # (see decode_ids), but the .smi input has only the base id. Strip it - # the same way reorder_sdf/count_output do, or every tautomer-derived - # molecule would be misreported as missing. - if "@taut" in name: - name = name.split("@taut")[0] - sdf_data.append(name) - sdf_data = list(set(sdf_data)) - - # Find molecules without 3D structures - bad: list[tuple[str, str]] = [] - for smiles_str, mol_id in smi_names: - if mol_id not in sdf_data: - bad.append((mol_id, smiles_str)) - - if len(bad) > 0: - logger.warning("The following SMILES has no 3D structure in the SDF file.") - logger.warning("ID, SMILES") - for mol_id, smiles_str in bad: - logger.warning(f"{mol_id} {smiles_str}") - else: - logger.info("Every SMILES has at least an 3D structure in the SDF file.") - - return bad - - -def find_ids_not_in_sdf(source_sdf: str, sdf: str) -> list[str]: - """Find molecule IDs from an SDF input that have no 3D structure in the output SDF. - - The SDF-input counterpart to :func:`find_smiles_not_in_sdf`. That function - reads its expected-IDs list from a ``.smi`` file, which does not exist when - the pipeline's input is itself an SDF file; this reads the same expected-IDs - list from the source SDF's ``_Name`` property instead, so SDF-input runs get - the same input/output reconciliation SMILES-input runs do. - - Args: - source_sdf: Path to the original input SDF file (pre-encoding IDs). - sdf: Path to the output SDF file (decoded IDs). - - Returns: - Input molecule IDs with no corresponding structure in ``sdf``. A source - record RDKit could not parse has no ``_Name`` to return, so it appears as - ``UNPARSEABLE_RECORD_ID`` filled in with its position -- it is a molecule - the user supplied and did not get back, and omitting it is what let a run - exit 0 having processed fewer molecules than its input contained. - - Example: - >>> missing = find_ids_not_in_sdf("input.sdf", "output.sdf") - >>> for mol_id in missing: - ... print(f"Failed: {mol_id}") - """ - # Find all input molecule IDs. - # - # A record RDKit cannot parse is reported, not skipped. `encode_ids` drops - # such a record with a warning so it never enters the run; this function then - # built its expected-ID list from the SAME file and skipped the SAME record, - # so the record was in neither `source_ids` nor the output, could not appear - # in `failures`, and `_exit_if_incomplete` saw `failed_count == 0`. The run - # printed a success summary and exited 0 having processed fewer molecules than - # the file contained -- exactly what this reconciliation exists to prevent - # (audit C7). It has no `_Name` to report, so it is named by position. - # - # Only the SDF path needed this. `encode_ids` reads `.smi` input with - # `on_malformed="raise"`, so a malformed SMILES line aborts the run with - # InputValidationError long before reconciliation. That the two input formats - # disagree on strictness -- `.smi` refuses the file, `.sdf` processes the rest - # -- is a real divergence, but unifying it changes behavior for large files - # and belongs with the other validation-consistency work, not here. - source_ids: list[str] = [] - for i, mol in enumerate(Chem.SDMolSupplier(source_sdf, removeHs=False)): - if mol is None: - # Not "Skipping ...", which is what this said while it was in fact - # skipping: the record is now counted as a failure, and a message - # claiming otherwise would be the same defect one layer up. - logger.warning( - "Input record at index %d could not be parsed; reporting it as a " - "molecule that produced no output.", - i, - ) - source_ids.append(UNPARSEABLE_RECORD_ID.format(index=i)) - continue - source_ids.append(mol.GetProp("_Name").strip()) - - # Get all molecule names from the output SDF - sdf_ids: set[str] = set() - mols = Chem.SDMolSupplier(sdf) - for i, mol in enumerate(mols): - if mol is None: - logger.warning("Skipping molecule at index %d: failed to parse", i) - continue - name = mol.GetProp("_Name") - if "@taut" in name: - name = name.split("@taut")[0] - sdf_ids.add(name) - - # Find molecules without 3D structures, preserving source order and - # de-duplicating (an id can appear once per tautomer/isomer conformer - # group in some callers, though not in the raw source SDF). - bad: list[str] = [] - seen: set[str] = set() - for mol_id in source_ids: - if mol_id not in sdf_ids and mol_id not in seen: - bad.append(mol_id) - seen.add(mol_id) - - if bad: - logger.warning("The following input IDs have no 3D structure in the SDF file.") - for mol_id in bad: - logger.warning(mol_id) - else: - logger.info("Every input molecule has at least one 3D structure in the SDF file.") - - return bad diff --git a/src/Auto3D/utils/geometry.py b/src/Auto3D/utils/geometry.py new file mode 100644 index 00000000..ca8c5018 --- /dev/null +++ b/src/Auto3D/utils/geometry.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python +"""Geometric measurements on a conformer's coordinates. + +Distances and RMSD only. Nothing here reads or writes a molecular property, +consults a force field, or decides whether a structure is acceptable -- those +belong to ``utils/connectivity.py``, ``Auto3D.clash_relief`` and +``Auto3D.filtering`` respectively. +""" +from __future__ import annotations + +import logging + +import numpy as np +from rdkit import Chem +from rdkit.Chem import rdMolAlign + +logger = logging.getLogger("auto3d") + +__all__ = ["min_pairwise_distance", "get_rmsd"] + + +def min_pairwise_distance(points: np.ndarray) -> float: + """Find the minimum pairwise distance among n points in 3D space. + + This function computes all pairwise distances between the provided points + and returns the minimum distance. It uses vectorized NumPy operations + for efficiency. + + Args: + points: A (n, 3) array representing the coordinates of n points + in 3D space. + + Returns: + The minimum pairwise distance among the n points. + + Example: + >>> import numpy as np + >>> points = np.array([[0, 0, 0], [1, 0, 0], [0, 2, 0]]) + >>> min_pairwise_distance(points) + 1.0 + """ + # Ensure input is a NumPy array with float32 type + points = points.astype(np.float32) + n = points.shape[0] + + # Guard for single atom or empty input + if n < 2: + # Single atom: no pairwise distance exists + return float('inf') + + # Expand dimensions of points to enable broadcasting + points_expanded = np.expand_dims(points, axis=1).repeat(n, axis=1) + + # Compute pairwise squared differences + diff_squared = (points_expanded - points_expanded.transpose(1, 0, 2)) ** 2 + + # Sum along the last dimension to get pairwise squared distances + pairwise_squared_distances = np.sum(diff_squared, axis=-1) + + # Find the minimum squared distance from upper triangle + upp_indices = np.triu_indices(n, 1) + upp_values = pairwise_squared_distances[upp_indices] + min_squared_distance = np.min(upp_values) + + # Return the square root of the minimum squared distance + return float(np.sqrt(min_squared_distance)) + + +def get_rmsd(mol1: Chem.Mol, mol2: Chem.Mol, remove_hs: bool = True) -> float: + """Calculate the RMSD between two molecular conformers. + + Uses RDKit's GetBestRMS function which finds the optimal alignment + between the two molecules before computing RMSD. + + Args: + mol1: First RDKit Mol object with a conformer. + mol2: Second RDKit Mol object with a conformer. + remove_hs: If True (default), remove hydrogens before RMSD calculation. + This speeds up the calculation and focuses on heavy atom positions. + + Returns: + The RMSD value in Angstroms. Returns ``float("inf")`` if alignment + fails (e.g., due to atom mismatch). An incomparable pair is treated as + "distinct" rather than "identical", which is the same convention used + by ``filter_unique``; a downstream ``rmsd < threshold`` check therefore + keeps the structure instead of dropping it as a false duplicate. + + Example: + >>> from rdkit import Chem + >>> from rdkit.Chem import AllChem + >>> mol1 = Chem.MolFromSmiles("CCO") + >>> mol1 = Chem.AddHs(mol1) + >>> AllChem.EmbedMolecule(mol1) + 0 + >>> mol2 = Chem.Mol(mol1) # Copy + >>> get_rmsd(mol1, mol2) + 0.0 + """ + try: + if remove_hs: + mol1_proc = Chem.RemoveHs(mol1) + mol2_proc = Chem.RemoveHs(mol2) + else: + mol1_proc = mol1 + mol2_proc = mol2 + # Temporary bug fix for https://github.com/rdkit/rdkit/issues/6826 + rmsd = rdMolAlign.GetBestRMS(mol1_proc, mol2_proc) + except RuntimeError: + # Incomparable pair: treat as distinct (inf), matching filter_unique. + rmsd = float("inf") + return float(rmsd) diff --git a/src/Auto3D/utils/logging_config.py b/src/Auto3D/utils/logging_config.py index b3c55a57..852cba5f 100644 --- a/src/Auto3D/utils/logging_config.py +++ b/src/Auto3D/utils/logging_config.py @@ -22,7 +22,7 @@ def get_logger(name: str) -> logging.Logger: A run's on-disk log (Auto3D.log) is fed separately, through a multiprocessing queue: ``Auto3D.workflow_workers`` attaches a ``QueueHandler`` onto BOTH this "Auto3D" tree and the lowercase "auto3d" - tree that ``Auto3D.workflow``, ``Auto3D.utils.chemistry`` and one warning + tree that ``Auto3D.workflow``, ``Auto3D.clash_relief`` and one warning in ``Auto3D.batch_opt.batchopt`` log through directly -- "auto3d" and "Auto3D" are case-distinct, unrelated sibling trees under root, so a warning from a logger returned here now reaches the run log too, without diff --git a/src/Auto3D/utils/molprops.py b/src/Auto3D/utils/molprops.py new file mode 100644 index 00000000..40d9f781 --- /dev/null +++ b/src/Auto3D/utils/molprops.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python +"""Scalar properties read straight off a molecular graph. + +Formal charge and the conformer budget: both are functions of the graph alone +(no coordinates, no force field, no energy), which is what separates them from +``utils/geometry.py`` and ``utils/connectivity.py``. +""" +from __future__ import annotations + +from rdkit import Chem +from rdkit.Chem import rdMolDescriptors, rdmolops + +from Auto3D.constants import ( + CONFORMER_MULTIPLIER, + CONFORMER_ROTATABLE_COEFF, + CONFORMER_ROTATABLE_EXP, + MAX_CONFORMERS_CAP, +) + +__all__ = ["calculate_conformer_count", "get_mol_charge"] + + +def calculate_conformer_count(mol: Chem.Mol) -> int: + """Calculate the number of conformers to generate for a molecule. + + Uses a formula based on the number of rotatable bonds, with a minimum + of the heavy atom count and a maximum cap. The result is floored at 1 so + a molecule never gets 0 conformers (which would silently drop tiny species + such as ``[H+]`` or a lone atom from the pipeline). + + Formula: min(max(1, num_heavy, 2 * 8.481 * (num_rotatable ** 1.642)), 1000) + Reference: https://doi.org/10.1021/acs.jctc.0c01213 + + Args: + mol: RDKit molecule object (with or without hydrogens). + + Returns: + Number of conformers to generate (always >= 1). + + Example: + >>> from rdkit import Chem + >>> mol = Chem.MolFromSmiles("CCCCCC") # hexane + >>> count = calculate_conformer_count(mol) + >>> 1 <= count <= 1000 + True + """ + num_rotatable = rdMolDescriptors.CalcNumRotatableBonds(mol) + num_heavy = sum(1 for atom in mol.GetAtoms() if atom.GetAtomicNum() > 1) + + formula_count = int( + CONFORMER_MULTIPLIER * CONFORMER_ROTATABLE_COEFF * + (num_rotatable ** CONFORMER_ROTATABLE_EXP) + ) + + # Floor at 1: a heavy-atom-free species (e.g. [H+]) or a single atom must + # still receive at least one conformer instead of being silently dropped. + return min(max(1, num_heavy, formula_count), MAX_CONFORMERS_CAP) + + +def get_mol_charge(mol: Chem.Mol) -> int: + """Get the formal charge of a molecule. + + Args: + mol: RDKit Mol object. + + Returns: + The total formal charge of the molecule. + + Example: + >>> from rdkit import Chem + >>> mol = Chem.MolFromSmiles("[NH4+]") + >>> get_mol_charge(mol) + 1 + """ + return rdmolops.GetFormalCharge(mol) diff --git a/src/Auto3D/utils/output_guard.py b/src/Auto3D/utils/output_guard.py new file mode 100644 index 00000000..0760452e --- /dev/null +++ b/src/Auto3D/utils/output_guard.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python +"""The two guards every Auto3D writer runs before it opens an output file. + +Split out of ``utils/validation.py`` so that a module which only needs to +refuse clobbering a file does not have to import that module -- and with it +``torch``, ``rdkit`` and (transitively, through the engine-name resolution) +the whole ``Auto3D.models`` tree. This module imports nothing but ``os`` and +``Auto3D.exceptions``, which is what lets the ``.smi``/``.sdf`` writers under +``utils/`` and the top-level ID/layout helpers gate their own output. + +``utils/validation.py`` still re-exports both names, because ``SPE``, +``ASE.thermo`` and ``cli.commands.properties`` import them from there. +""" +from __future__ import annotations + +import os + +from Auto3D.exceptions import ConfigurationError + +__all__ = ["check_output_not_input", "check_output_overwrite"] + + +def check_output_not_input(path: str, out_path: str | None) -> None: + """Refuse to write the output over the input file. + + ``auto3d energy mols.sdf -o mols.sdf`` used to open ``mols.sdf`` for + writing while the run was still reading from it, so the user's input was + destroyed -- and, if the run then failed part-way, replaced by a truncated + file with no surviving copy of either the input or the result (C14). The + Phase 6 tmp+``os.replace`` staging fixes the *crash* half of C14 (a failed + rewrite no longer leaves a partial file), but it cannot fix this half: a + successful same-file run still deliberately overwrites the input, and no + amount of atomicity brings the original back. + + Single source of truth for that policy, in the same spirit as + ``check_gpu_requested`` and ``check_engine_supports_molecules``: + ``calc_spe``, ``opt_geometry`` and ``calc_thermo`` each take an output path + directly and never go through ``check_input``/``check_valid_configuration``, + so all three call this function rather than carrying three copies of the + test that would drift apart. The ``auto3d energy``/``optimize``/``thermo`` + CLI commands pass ``--output`` straight through to those functions, so they + are covered by the same call. + + Two comparisons, because neither alone is sufficient: + + ``os.path.samefile`` is the authoritative test -- it compares ``st_dev`` and + ``st_ino``, so it catches the two cases string/``realpath`` comparison + misses entirely. A **hardlink** (``cp -l mols.sdf results.sdf``) is one file + under two names with two distinct real paths, so ``realpath`` compares them + unequal and writing to either destroys the other. A **case-insensitive + filesystem** (macOS APFS/HFS+, Windows NTFS -- both supported platforms) + resolves ``Mols.sdf`` and ``mols.sdf`` to one file whose real paths differ + only in case. Both defeat ``realpath`` equality; ``samefile`` sees through + both because the kernel already told it they are the same inode. + + ``samefile`` requires both paths to exist, and in the normal case the output + does not yet -- so it is guarded by ``os.path.exists`` and the ``realpath`` + comparison is kept as the fallback. That fallback is what catches the common + spellings (``mols.sdf`` vs ``./mols.sdf`` vs an absolute path vs a symlink) + when the output file has not been created yet, which ``samefile`` cannot + answer at all. + + Args: + path: The input file the caller will read. + out_path: The requested output path, or None to use the default + (which is derived from `path` and never equals it). + + Raises: + ConfigurationError: `out_path` names the same file as `path`. + """ + if out_path is None: + return + + same = os.path.realpath(path) == os.path.realpath(out_path) + if not same and os.path.exists(path) and os.path.exists(out_path): + try: + same = os.path.samefile(path, out_path) + except OSError: + # A path that vanished between exists() and samefile(), or that + # cannot be stat'd. Fall back to the realpath verdict rather than + # failing the run on a check that is itself best-effort. + pass + + if same: + raise ConfigurationError( + f"Output path {out_path!r} is the same file as the input {path!r}. " + "Auto3D would overwrite your input; pass a different output path." + ) + + +def check_output_overwrite(out_path: str | os.PathLike[str] | None, overwrite: bool) -> None: + """Refuse to write over a file that already exists. + + ``auto3d energy junk.sdf --no-gpu -o precious.sdf`` used to exit 0, print + "Wrote precious.sdf", and leave ``precious.sdf`` at **0 bytes**: every + writer below opens ``Chem.SDWriter(outpath)``, which truncates on open, + and ``calc_spe`` takes an early-return branch that opens the writer and + writes nothing when every record in the input fails to parse. + + Be precise about *when* the destruction happened, because it is not what + "truncates on open" suggests: all four writers open their output only + after the compute is finished (``SPE.py:161``, ``ASE/thermo.py:878``, + ``batch_opt/batchopt.py:323`` for ``opt_geometry``, ``ranking.py:287``), + so a run that failed part-way left the user's file untouched. What + destroyed it was a run that *succeeded*, or -- for the 0-byte case above + -- one that had nothing to write. This guard exists because both of those + are silent: nothing warned that the path was occupied. ``auto3d config init`` has refused to + clobber an existing file since it shipped; the calculators did not. + + Single source of truth for that policy, in the same spirit as + ``check_output_not_input`` directly above: ``calc_spe``, ``opt_geometry``, + ``calc_thermo`` and ``ConformerRanker`` each resolve their own output path + and would otherwise each carry their own copy of this test, which is how + four copies drift apart. ``auto3d tautomers`` derives its output name + inside the pipeline and honors ``-o`` with a ``shutil.move``, so its CLI + wrapper calls this function itself before the pipeline runs. + + This is a *distinct* guard from ``check_output_not_input``, not a + generalization of it: that one refuses ``out_path`` naming the input even + when ``--force`` is passed (there is no recovering an input you overwrote + with a filtered subset of itself), while this one is a consent gate the + user can lift. Both run; neither subsumes the other. + + The check is on the *resolved* output path, so it covers the default + derived name (``mols_AIMNET_E.sdf``) exactly as it covers an explicit + ``-o``. A second ``auto3d energy mols.sdf`` therefore stops rather than + silently replacing the first run's results. + + ``os.path.exists`` follows symlinks, which is the behavior wanted here: a + symlink pointing at a real file is a file the write would destroy. A + dangling symlink reports False and is overwritten, matching what the + writer would do anyway. + + Args: + out_path: The resolved path the caller is about to write, or None + when the caller has nothing to write. + overwrite: True to allow clobbering an existing file (``--force`` on + the CLI, ``overwrite=True`` in the Python API). + + Raises: + ConfigurationError: `out_path` exists and `overwrite` is False. + """ + if out_path is None or overwrite: + return + + if os.path.exists(out_path): + raise ConfigurationError( + f"{out_path} already exists. Pass --force/-f to overwrite, or " + "choose a different -o path. (Python API: pass overwrite=True.)", + # No hint: the message above already states both ways out, and + # ConfigurationError's class hint ("run auto3d config init") has + # nothing to do with an -o collision. "" suppresses it; None + # would have meant "unset" and let the class hint through. + hint="", + ) diff --git a/src/Auto3D/utils/reconciliation.py b/src/Auto3D/utils/reconciliation.py new file mode 100644 index 00000000..621f5f38 --- /dev/null +++ b/src/Auto3D/utils/reconciliation.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python +"""Reconciling what a run was given against what it produced. + +One question, asked once per input format: which molecules did the user hand +Auto3D that came back with no 3D structure? A run that quietly processes fewer +molecules than its input contained and still exits 0 is the defect these +helpers exist to prevent, so both of them report a record they could not even +parse rather than skipping it. +""" +from __future__ import annotations + +from rdkit import Chem + +from Auto3D.utils.logging_config import get_logger +from Auto3D.utils.smi_io import iter_smi_records + +logger = get_logger(__name__) + +#: Stand-in ID for an input record Auto3D could not parse, used by the +#: input-vs-output reconciliation. Such a record has no ``_Name`` to report, so it +#: is identified by its position in the source file. The angle brackets make it +#: unmistakable as a placeholder rather than a molecule name a user chose -- an ID +#: read from a file can be anything, including ``record 3``, but not with these. +UNPARSEABLE_RECORD_ID = "" + + +def find_smiles_not_in_sdf(smi: str, sdf: str) -> list[tuple[str, str]]: + """Find SMILES that failed to generate 3D conformers. + + Compares a SMILES input file against an SDF output file to identify + molecules that did not successfully generate 3D structures. + + Args: + smi: Path to input SMILES file. + sdf: Path to output SDF file. + + Returns: + List of (id, smiles) tuples for molecules not in SDF. + + Example: + >>> bad = find_smiles_not_in_sdf("input.smi", "output.sdf") + >>> for mol_id, smiles in bad: + ... print(f"Failed: {mol_id}") + """ + # Find all SMILES ids + smi_names: list[tuple[str, str]] = [] + for _line_no, smiles_str, mol_id in iter_smi_records(smi, on_malformed="skip"): + smi_names.append((smiles_str, mol_id)) + + # Get all molecule names from SDF + sdf_data: list[str] = [] + mols = Chem.SDMolSupplier(sdf) + for i, mol in enumerate(mols): + if mol is None: + logger.warning("Skipping molecule at index %d: failed to parse", i) + continue + name = mol.GetProp("_Name") + # decode_ids keeps a "@tautN" suffix on tautomer-enumerated conformers + # (see Auto3D.id_mapping.decode_ids), but the .smi input has only the + # base id. Strip it the same way reorder_sdf/count_output do, or every + # tautomer-derived molecule would be misreported as missing. + if "@taut" in name: + name = name.split("@taut")[0] + sdf_data.append(name) + sdf_data = list(set(sdf_data)) + + # Find molecules without 3D structures + bad: list[tuple[str, str]] = [] + for smiles_str, mol_id in smi_names: + if mol_id not in sdf_data: + bad.append((mol_id, smiles_str)) + + if len(bad) > 0: + logger.warning("The following SMILES has no 3D structure in the SDF file.") + logger.warning("ID, SMILES") + for mol_id, smiles_str in bad: + logger.warning(f"{mol_id} {smiles_str}") + else: + logger.info("Every SMILES has at least an 3D structure in the SDF file.") + + return bad + + +def find_ids_not_in_sdf(source_sdf: str, sdf: str) -> list[str]: + """Find molecule IDs from an SDF input that have no 3D structure in the output SDF. + + The SDF-input counterpart to :func:`find_smiles_not_in_sdf`. That function + reads its expected-IDs list from a ``.smi`` file, which does not exist when + the pipeline's input is itself an SDF file; this reads the same expected-IDs + list from the source SDF's ``_Name`` property instead, so SDF-input runs get + the same input/output reconciliation SMILES-input runs do. + + Args: + source_sdf: Path to the original input SDF file (pre-encoding IDs). + sdf: Path to the output SDF file (decoded IDs). + + Returns: + Input molecule IDs with no corresponding structure in ``sdf``. A source + record RDKit could not parse has no ``_Name`` to return, so it appears as + ``UNPARSEABLE_RECORD_ID`` filled in with its position -- it is a molecule + the user supplied and did not get back, and omitting it is what let a run + exit 0 having processed fewer molecules than its input contained. + + Example: + >>> missing = find_ids_not_in_sdf("input.sdf", "output.sdf") + >>> for mol_id in missing: + ... print(f"Failed: {mol_id}") + """ + # Find all input molecule IDs. + # + # A record RDKit cannot parse is reported, not skipped. `encode_ids` drops + # such a record with a warning so it never enters the run; this function then + # built its expected-ID list from the SAME file and skipped the SAME record, + # so the record was in neither `source_ids` nor the output, could not appear + # in `failures`, and `_exit_if_incomplete` saw `failed_count == 0`. The run + # printed a success summary and exited 0 having processed fewer molecules than + # the file contained -- exactly what this reconciliation exists to prevent + # (audit C7). It has no `_Name` to report, so it is named by position. + # + # Only the SDF path needed this. `encode_ids` reads `.smi` input with + # `on_malformed="raise"`, so a malformed SMILES line aborts the run with + # InputValidationError long before reconciliation. That the two input formats + # disagree on strictness -- `.smi` refuses the file, `.sdf` processes the rest + # -- is a real divergence, but unifying it changes behavior for large files + # and belongs with the other validation-consistency work, not here. + source_ids: list[str] = [] + for i, mol in enumerate(Chem.SDMolSupplier(source_sdf, removeHs=False)): + if mol is None: + # Not "Skipping ...", which is what this said while it was in fact + # skipping: the record is now counted as a failure, and a message + # claiming otherwise would be the same defect one layer up. + logger.warning( + "Input record at index %d could not be parsed; reporting it as a " + "molecule that produced no output.", + i, + ) + source_ids.append(UNPARSEABLE_RECORD_ID.format(index=i)) + continue + source_ids.append(mol.GetProp("_Name").strip()) + + # Get all molecule names from the output SDF + sdf_ids: set[str] = set() + mols = Chem.SDMolSupplier(sdf) + for i, mol in enumerate(mols): + if mol is None: + logger.warning("Skipping molecule at index %d: failed to parse", i) + continue + name = mol.GetProp("_Name") + if "@taut" in name: + name = name.split("@taut")[0] + sdf_ids.add(name) + + # Find molecules without 3D structures, preserving source order and + # de-duplicating (an id can appear once per tautomer/isomer conformer + # group in some callers, though not in the raw source SDF). + bad: list[str] = [] + seen: set[str] = set() + for mol_id in source_ids: + if mol_id not in sdf_ids and mol_id not in seen: + bad.append(mol_id) + seen.add(mol_id) + + if bad: + logger.warning("The following input IDs have no 3D structure in the SDF file.") + for mol_id in bad: + logger.warning(mol_id) + else: + logger.info("Every input molecule has at least one 3D structure in the SDF file.") + + return bad diff --git a/src/Auto3D/utils/sdf_io.py b/src/Auto3D/utils/sdf_io.py new file mode 100644 index 00000000..16499d76 --- /dev/null +++ b/src/Auto3D/utils/sdf_io.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python +"""Reading, splitting, counting and reordering SDF files. + +Structural SDF file handling only: nothing here knows what an Auto3D energy or +convergence flag means (``utils/energy.py`` and ``utils/convergence.py`` own +those), and nothing here decides pipeline layout (``Auto3D.job_layout``) or ID +policy (``Auto3D.id_mapping``). +""" +from __future__ import annotations + +from collections import defaultdict +from pathlib import Path + +from rdkit import Chem + +from Auto3D.utils.atomic_io import atomic_write_path +from Auto3D.utils.logging_config import get_logger +from Auto3D.utils.smi_io import iter_smi_records + +logger = get_logger(__name__) + + +def guess_file_type(filename: str) -> str: + """Return the file extension for a given filename. + + Determines the file type based on the extension of the provided filename. + The extension is returned without the leading dot. + + Args: + filename: Path or filename to analyze. + + Returns: + The file extension without the leading dot (e.g., 'smi', 'sdf', 'xyz'). + + Example: + >>> guess_file_type("molecules.sdf") + 'sdf' + >>> guess_file_type("/path/to/input.smi") + 'smi' + >>> guess_file_type("file.mol2") + 'mol2' + """ + return Path(filename).suffix[1:] + + +def SDF2chunks(sdf: str) -> list[list[str]]: + """Split an SDF file into chunks, one per molecule. + + Reads an SDF file and splits it into a list of chunks, where each chunk + contains the lines of a single molecule as they appear in the original file. + + Args: + sdf: Path to the input SDF file. + + Returns: + List of chunks, where each chunk is a list of strings (lines) + representing one molecule including the '$$$$' terminator. + + Example: + >>> chunks = SDF2chunks("molecules.sdf") + >>> len(chunks) # Number of molecules + 10 + >>> chunks[0][-1].strip() # Last line of first molecule + '$$$$' + """ + chunks: list[list[str]] = [] + with open(sdf) as f: + data = f.readlines() + chunk: list[str] = [] + for line in data: + if line.strip() == "$$$$": + chunk.append(line) + chunks.append(chunk) + chunk = [] + else: + chunk.append(line) + # A final record lacking the '$$$$' terminator leaves residual lines in + # `chunk`. Preserve it as the last chunk rather than silently dropping it. + if any(line.strip() for line in chunk): + logger.warning( + "SDF file %s ends without a '$$$$' terminator; " + "keeping the trailing record as a final chunk.", + sdf, + ) + chunks.append(chunk) + return chunks + + +def reorder_sdf(sdf: str, source: str) -> list[Chem.Mol]: + """Reorder conformers in an SDF file to match the input source file order. + + Reads the order of molecule IDs from the source file and rewrites the SDF + file with conformers ordered to match. This ensures consistent output + ordering regardless of processing order. + + Args: + sdf: Path to the SDF file to reorder (will be overwritten). + source: Path to the source .smi or .sdf file defining the desired order. + + Returns: + List of RDKit Mol objects in the reordered sequence. + + Note: + - For tautomer conformers (containing '@taut' in ID), the base ID + is extracted for ordering purposes. + - If the source format is unsupported, prints a message and returns None. + - Molecules whose id is not present in ``source`` are appended at the + end (not dropped), so no data is lost. + - Duplicate source ids are de-duplicated: each id's molecules are + written once, so the returned list may be shorter than the input if + source ids repeat. + + Example: + >>> ordered_mols = reorder_sdf("output_3d.sdf", "input.smi") + >>> len(ordered_mols) + 10 + """ + # convert smi/sdf to a list of ids with correct order + ids: list[str] = [] + format = guess_file_type(source) + if format == "smi": + for _line_no, _smiles, mol_id in iter_smi_records(source, on_malformed="skip"): + ids.append(mol_id) + elif format == "sdf": + supp = Chem.SDMolSupplier(source, removeHs=False) + for i, mol in enumerate(supp): + if mol is None: + logger.warning("Skipping molecule at index %d: failed to parse", i) + continue + ids.append(mol.GetProp("_Name")) + else: + logger.warning("Unsupported file format: %s" % format) + return None # type: ignore + + # convert sdf to a Dict[id, List[mols]], preserving discovery order so any + # molecule whose id is not in `source` can still be appended (no data loss). + id_mols: dict[str, list[Chem.Mol]] = defaultdict(lambda: []) + discovery_order: list[str] = [] + supp = Chem.SDMolSupplier(sdf, removeHs=False) + for i, mol in enumerate(supp): + if mol is None: + logger.warning("Skipping molecule at index %d: failed to parse", i) + continue + id = mol.GetProp("_Name") + if "@taut" in id: + id = id.split("@taut")[0] + if id not in id_mols: + discovery_order.append(id) + id_mols[id].append(mol) + + # Release the RDKit supplier's file handle before overwriting `sdf`. + # On Windows an open handle makes the later os.replace() fail with + # "Access is denied" (WinError 5); on POSIX the replace would succeed. + del supp + + # Order: ids present in `source` first (in source order), then any + # unmatched molecules appended in their original order so nothing is lost. + source_id_set = set(ids) + ordered_ids = list(ids) + for id in discovery_order: + if id not in source_id_set: + logger.warning( + "Molecule id %r in %s is not present in source %s; " + "appending it at the end to avoid data loss.", + id, + sdf, + source, + ) + ordered_ids.append(id) + + # Write the mols in the correct order to a sibling temp file, then + # atomically replace the original only on success (crash-safe in-place + # overwrite). `atomic_write_path` owns that staging for all three of + # Auto3D's in-place rewrites; this one used to do it by hand through a + # predictable `.reorder.tmp` and without copying `sdf`'s permission + # bits, so a 0600 file came back at whatever the umask allows. + ordered_mols: list[Chem.Mol] = [] + written_ids: set[str] = set() + with atomic_write_path(sdf, suffix=".sdf") as tmp_path, Chem.SDWriter(tmp_path) as f: + for id in ordered_ids: + if id in written_ids: + continue + written_ids.add(id) + mols = id_mols[id] + if len(mols) >= 1: + ordered_mols.extend(mols) + for mol in mols: + f.write(mol) + return ordered_mols + + +def count_sdf(sdf: str) -> int: + """Count the number of molecules in an SDF file. + + Args: + sdf: Path to the SDF file. + + Returns: + Number of molecules in the file. + + Example: + >>> count_sdf("molecules.sdf") + 10 + """ + mols = Chem.SDMolSupplier(sdf) + return len([mol for mol in mols if mol is not None]) diff --git a/src/Auto3D/utils/smi_io.py b/src/Auto3D/utils/smi_io.py new file mode 100644 index 00000000..4781705d --- /dev/null +++ b/src/Auto3D/utils/smi_io.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python +"""Reading and writing ``.smi`` files. + +The one place that knows what a line of an Auto3D ``.smi`` file looks like: +``SMILES ID`` with any further whitespace-separated columns ignored, blank +lines skipped and ``#`` comment lines skipped. Everything that consumes or +produces that format goes through :func:`iter_smi_records` or one of the +writers here, so ``auto3d validate``, the run pipeline, the isomer/tautomer +engines and the input/output reconciliation cannot drift apart on it. + +Deliberately free of ``torch``: these writers are reached from +``Auto3D.utils`` leaves and from the ID-encoding step, and pulling the model +tree in through a validation import would make an ordinary ``.smi`` write cost +the whole ``Auto3D.models`` package. The overwrite guards live in +``utils/output_guard.py`` for the same reason. +""" +from __future__ import annotations + +import collections + +from rdkit import Chem +from rdkit.Chem import inchi + +from Auto3D.exceptions import InputValidationError +from Auto3D.utils.logging_config import get_logger + +logger = get_logger(__name__) + + +def iter_smi_records(path, *, on_malformed="skip"): + """Yield (line_no, smiles, mol_id) for each non-blank, non-comment line of + a .smi file. + + A line is 'SMILES ID [extra columns ignored]'. Blank/whitespace-only lines + are skipped, as are lines whose first non-whitespace character is '#' + (comments) -- matching cli.commands.validate.validate_smiles_file, so + `auto3d validate` and every consumer of this function (encode_ids and so + the whole run pipeline, plus the isomer/tautomer engines and the + input/output reconciliation helpers) agree on what a comment line is + (M25). A real SMILES token can never start with '#' (it is a bond symbol + between two atoms, never a leading character), so this cannot misclassify + a legitimate SMILES as a comment. on_malformed controls lines with fewer + than 2 whitespace tokens: + - "skip": log a warning and skip the line (lenient; default) + - "raise": raise InputValidationError naming the line + + Args: + path: Path to the input .smi file. + on_malformed: How to handle lines with fewer than 2 tokens + ("skip" or "raise"). + + Yields: + Tuples of (line_no, smiles, mol_id) where line_no is 1-based. Any extra + whitespace-separated columns beyond the ID are intentionally ignored. + + Raises: + InputValidationError: If on_malformed == "raise" and a non-blank, + non-comment line has fewer than 2 whitespace tokens. + ValueError: If on_malformed is not "skip" or "raise". + """ + if on_malformed not in ("skip", "raise"): + raise ValueError( + f"on_malformed must be 'skip' or 'raise', got: {on_malformed!r}" + ) + with open(path) as f: + data = f.readlines() + for line_no, line in enumerate(data, start=1): + stripped = line.strip() + if not stripped: + continue + if stripped.startswith("#"): + continue + parts = stripped.split() + if len(parts) < 2: + if on_malformed == "raise": + raise InputValidationError( + f"Line {line_no} is missing a molecule ID " + f"(expected 'SMILES ID'): {line.strip()!r}" + ) + logger.warning( + f"Skipping molecule at line {line_no}: failed to parse " + f"(need 'SMILES ID', got: {line.strip()!r})" + ) + continue + # Lenient parsing: ignore any extra whitespace-separated columns. + yield line_no, parts[0], parts[1] + + +def smiles2smi(smiles: list[str], path: str) -> str: + """Convert a list of SMILES strings to a .smi file with InChIKey IDs. + + Each SMILES string is converted to a molecule, and its InChIKey is computed + to serve as a unique identifier. The output file contains one molecule per + line in the format: "SMILES InChIKey". + + Args: + smiles: List of SMILES strings to convert. + path: Output file path for the .smi file. + + Returns: + The output file path. + + Example: + >>> smiles2smi(["CCO", "CCC"], "molecules.smi") + 'molecules.smi' + # File content: + # CCO LFQSCWFLJHTTHZ-UHFFFAOYSA-N + # CCC ATUOYWHBWRKTHZ-UHFFFAOYSA-N + """ + lines = [] + seen_ids: dict[str, int] = {} + for idx, smi in enumerate(smiles): + mol = Chem.MolFromSmiles(smi) + if mol is None: + raise InputValidationError( + f"Invalid SMILES at index {idx}: {smi!r} could not be parsed " + "by RDKit." + ) + inchikey = inchi.MolToInchiKey(mol) + # Distinct inputs can share a standard InChIKey (e.g. tautomers the + # standard InChIKey conflates, or the same molecule written two ways). + # The InChIKey is used as the molecule's unique ID downstream, and + # reorder_sdf collapses duplicate IDs -- so a colliding input would be + # silently dropped. Disambiguate by suffixing repeats (_2, _3, ...) so + # every input keeps its own conformers. The suffix stays a single + # whitespace-delimited token and round-trips through enumeration. + count = seen_ids.get(inchikey, 0) + 1 + seen_ids[inchikey] = count + mol_id = inchikey if count == 1 else f"{inchikey}_{count}" + if count > 1: + logger.info( + "Input SMILES %r shares InChIKey %s with an earlier input; " + "assigning disambiguated id %s so it is not dropped.", + smi, + inchikey, + mol_id, + ) + lines.append(f"{smi} {mol_id}\n") + + with open(path, "w+") as f: + for line in lines: + f.write(line) + + return path + + +def hash_enumerated_smi_IDs(smi: str, out: str) -> None: + """Write all SMILES with hashed IDs into a new file. + + Reads a SMILES file, sorts entries by ID, handles duplicate IDs by appending + '_0' suffix, and writes the result to the output file. + + Args: + smi: Path to the input .smi file. + out: Path for the output .smi file with sorted/hashed IDs. + + Returns: + None. Writes the result to the output file. + + Example: + >>> hash_enumerated_smi_IDs("input.smi", "output.smi") + """ + dict0: dict[str, str] = {} + for _line_no, smiles, id in iter_smi_records(smi, on_malformed="skip"): + while id in dict0: + id += "_0" + dict0[id] = smiles + + dict0 = collections.OrderedDict(sorted(dict0.items())) + + with open(out, "w+") as f: + for id, smiles in dict0.items(): + molecule = smiles.strip() + " " + id.strip() + "\n" + f.write(molecule) + + +def hash_taut_smi(smi: str, out: str) -> None: + """Write all SMILES with hashed IDs for tautomers. + + Reads a SMILES file and appends '@tautN' suffix to IDs where N is + an incrementing counter, ensuring unique tautomer identifiers. + + Args: + smi: Path to the input .smi file. + out: Path for the output .smi file with tautomer IDs. + + Returns: + None. Writes the result to the output file. + + Example: + >>> hash_taut_smi("input.smi", "tautomers.smi") + """ + dict0: dict[str, str] = {} + for _line_no, smiles, id in iter_smi_records(smi, on_malformed="skip"): + c = 1 + id_ = id + while ("taut" not in id_) or (id_ in dict0): + id_ = id + f"@taut{c}" + c += 1 + dict0[id_] = smiles + + dict0 = collections.OrderedDict(sorted(dict0.items())) + + with open(out, "w+") as f: + for id, smiles in dict0.items(): + molecule = smiles.strip() + " " + id.strip() + "\n" + f.write(molecule) + + +def combine_smi(smies: list[str], out: str) -> None: + """Combine multiple SMILES files into a single file. + + Reads all input SMILES files, removes duplicates, and writes the + combined unique entries to the output file. + + Args: + smies: List of paths to input .smi files. + out: Path for the combined output .smi file. + + Returns: + None. Writes the combined result to the output file. + + Example: + >>> combine_smi(["file1.smi", "file2.smi"], "combined.smi") + """ + data: list[str] = [] + for smi in smies: + with open(smi) as f: + datai = f.readlines() + data += datai + # Order-preserving dedup: list(set(...)) randomizes line order across runs + # (hash seed), making the combined output non-deterministic. dict.fromkeys + # keeps first-seen order while removing exact duplicates. + data = list(dict.fromkeys(data)) + with open(out, "w+") as f2: + for line in data: + if not line.isspace(): + f2.write(line.strip() + "\n") diff --git a/src/Auto3D/utils/stereochemistry.py b/src/Auto3D/utils/stereochemistry.py index 2ac0393c..f6d7e2af 100644 --- a/src/Auto3D/utils/stereochemistry.py +++ b/src/Auto3D/utils/stereochemistry.py @@ -10,15 +10,13 @@ from __future__ import annotations import math -import os import re -import stat -import tempfile from collections import OrderedDict, defaultdict from rdkit import Chem from rdkit.Chem.rdMolDescriptors import CalcNumAtomStereoCenters +from Auto3D.utils.atomic_io import atomic_write_path from Auto3D.utils.logging_config import get_logger logger = get_logger(__name__) @@ -263,7 +261,7 @@ def remove_enantiomers(inpath: str, out: str) -> dict[str, list[str]]: # Strip only the trailing isomer-index component write_enumerated_smi # appends (rsplit, maxsplit=1), not everything after the first # underscore: an id like "KEY_2" -- smiles2smi's disambiguation of a - # duplicate InChIKey (utils/file_ops.py), kept distinct from "KEY" + # duplicate InChIKey (utils/smi_io.py), kept distinct from "KEY" # specifically so it is not dropped -- must survive this grouping # intact, or it silently merges back onto "KEY" here before ranking # ever sees it (M17). @@ -560,42 +558,19 @@ def amend_configuration_w(smi: str) -> None: Note: The rewrite is staged through a sibling temp file and moved into place - with ``os.replace`` (atomic on POSIX and Windows). Opening ``smi`` for - writing directly would truncate it, so a failure partway through the - loop below would destroy the input this function just read and leave - nothing to recover from (C14). + with ``os.replace`` (atomic on POSIX and Windows) by + :func:`Auto3D.utils.atomic_io.atomic_write_path`, which owns that + staging for all three of Auto3D's in-place rewrites. Opening ``smi`` + for writing directly would truncate it, so a failure partway through + the loop below would destroy the input this function just read and + leave nothing to recover from (C14). """ dct = amend_configuration(smi) - # Same directory as the target: os.replace raises OSError across - # filesystems. mkstemp creates the file 0600, so copy the original's - # permission bits over rather than silently tightening the user's file. - # realpath, not abspath: abspath collapses ".." lexically, so a path like - # /scratch/link/../in.smi (link -> another mount) would stage the temp file - # on the wrong filesystem and os.replace would fail with EXDEV. Only the - # PARENT is resolved -- os.replace acts on the final component itself. - directory = os.path.realpath(os.path.dirname(os.path.abspath(smi))) - fd, tmp_path = tempfile.mkstemp(suffix=".smi", dir=directory) - os.close(fd) - try: - os.chmod(tmp_path, stat.S_IMODE(os.stat(smi).st_mode)) - except OSError: # pragma: no cover - best effort; never block the rewrite - pass - - try: - with open(tmp_path, "w") as f: - for key in dct.keys(): - val = dct[key] - for i, smi_str in enumerate(val): - idx = str(key).strip() + "_" + str(i + 1) - line = smi_str + " " + idx + "\n" - f.write(line) - os.replace(tmp_path, smi) - except BaseException: - # BaseException, not Exception: a KeyboardInterrupt mid-write must not - # leave a stray .smi beside the file being amended. - try: - os.unlink(tmp_path) - except OSError: - pass - raise + with atomic_write_path(smi, suffix=".smi") as tmp_path, open(tmp_path, "w") as f: + for key in dct.keys(): + val = dct[key] + for i, smi_str in enumerate(val): + idx = str(key).strip() + "_" + str(i + 1) + line = smi_str + " " + idx + "\n" + f.write(line) diff --git a/src/Auto3D/utils/validation.py b/src/Auto3D/utils/validation.py index 358c19c9..af2656a9 100644 --- a/src/Auto3D/utils/validation.py +++ b/src/Auto3D/utils/validation.py @@ -22,6 +22,16 @@ ModelLoadError, ) from Auto3D.utils.logging_config import get_logger + +# The two output guards moved to the leaf module ``utils/output_guard.py`` so +# that a writer needing only them does not import this module's torch. They are +# re-exported here because ``SPE``, ``ASE.thermo`` and +# ``cli.commands.properties`` import them from this path; new call sites should +# name ``Auto3D.utils.output_guard`` directly. +from Auto3D.utils.output_guard import ( # noqa: F401 + check_output_not_input, + check_output_overwrite, +) from Auto3D.utils.stereochemistry import count_unspecified_stereo if TYPE_CHECKING: @@ -142,141 +152,6 @@ def check_engine_supports_molecules( ) -def check_output_not_input(path: str, out_path: str | None) -> None: - """Refuse to write the output over the input file. - - ``auto3d energy mols.sdf -o mols.sdf`` used to open ``mols.sdf`` for - writing while the run was still reading from it, so the user's input was - destroyed -- and, if the run then failed part-way, replaced by a truncated - file with no surviving copy of either the input or the result (C14). The - Phase 6 tmp+``os.replace`` staging fixes the *crash* half of C14 (a failed - rewrite no longer leaves a partial file), but it cannot fix this half: a - successful same-file run still deliberately overwrites the input, and no - amount of atomicity brings the original back. - - Single source of truth for that policy, in the same spirit as - ``check_gpu_requested`` and ``check_engine_supports_molecules``: - ``calc_spe``, ``opt_geometry`` and ``calc_thermo`` each take an output path - directly and never go through ``check_input``/``check_valid_configuration``, - so all three call this function rather than carrying three copies of the - test that would drift apart. The ``auto3d energy``/``optimize``/``thermo`` - CLI commands pass ``--output`` straight through to those functions, so they - are covered by the same call. - - Two comparisons, because neither alone is sufficient: - - ``os.path.samefile`` is the authoritative test -- it compares ``st_dev`` and - ``st_ino``, so it catches the two cases string/``realpath`` comparison - misses entirely. A **hardlink** (``cp -l mols.sdf results.sdf``) is one file - under two names with two distinct real paths, so ``realpath`` compares them - unequal and writing to either destroys the other. A **case-insensitive - filesystem** (macOS APFS/HFS+, Windows NTFS -- both supported platforms) - resolves ``Mols.sdf`` and ``mols.sdf`` to one file whose real paths differ - only in case. Both defeat ``realpath`` equality; ``samefile`` sees through - both because the kernel already told it they are the same inode. - - ``samefile`` requires both paths to exist, and in the normal case the output - does not yet -- so it is guarded by ``os.path.exists`` and the ``realpath`` - comparison is kept as the fallback. That fallback is what catches the common - spellings (``mols.sdf`` vs ``./mols.sdf`` vs an absolute path vs a symlink) - when the output file has not been created yet, which ``samefile`` cannot - answer at all. - - Args: - path: The input file the caller will read. - out_path: The requested output path, or None to use the default - (which is derived from `path` and never equals it). - - Raises: - ConfigurationError: `out_path` names the same file as `path`. - """ - if out_path is None: - return - - same = os.path.realpath(path) == os.path.realpath(out_path) - if not same and os.path.exists(path) and os.path.exists(out_path): - try: - same = os.path.samefile(path, out_path) - except OSError: - # A path that vanished between exists() and samefile(), or that - # cannot be stat'd. Fall back to the realpath verdict rather than - # failing the run on a check that is itself best-effort. - pass - - if same: - raise ConfigurationError( - f"Output path {out_path!r} is the same file as the input {path!r}. " - "Auto3D would overwrite your input; pass a different output path." - ) - - -def check_output_overwrite(out_path: str | os.PathLike[str] | None, overwrite: bool) -> None: - """Refuse to write over a file that already exists. - - ``auto3d energy junk.sdf --no-gpu -o precious.sdf`` used to exit 0, print - "Wrote precious.sdf", and leave ``precious.sdf`` at **0 bytes**: every - writer below opens ``Chem.SDWriter(outpath)``, which truncates on open, - and ``calc_spe`` takes an early-return branch that opens the writer and - writes nothing when every record in the input fails to parse. - - Be precise about *when* the destruction happened, because it is not what - "truncates on open" suggests: all four writers open their output only - after the compute is finished (``SPE.py:161``, ``ASE/thermo.py:878``, - ``batch_opt/batchopt.py:323`` for ``opt_geometry``, ``ranking.py:287``), - so a run that failed part-way left the user's file untouched. What - destroyed it was a run that *succeeded*, or -- for the 0-byte case above - -- one that had nothing to write. This guard exists because both of those - are silent: nothing warned that the path was occupied. ``auto3d config init`` has refused to - clobber an existing file since it shipped; the calculators did not. - - Single source of truth for that policy, in the same spirit as - ``check_output_not_input`` directly above: ``calc_spe``, ``opt_geometry``, - ``calc_thermo`` and ``ConformerRanker`` each resolve their own output path - and would otherwise each carry their own copy of this test, which is how - four copies drift apart. ``auto3d tautomers`` derives its output name - inside the pipeline and honors ``-o`` with a ``shutil.move``, so its CLI - wrapper calls this function itself before the pipeline runs. - - This is a *distinct* guard from ``check_output_not_input``, not a - generalization of it: that one refuses ``out_path`` naming the input even - when ``--force`` is passed (there is no recovering an input you overwrote - with a filtered subset of itself), while this one is a consent gate the - user can lift. Both run; neither subsumes the other. - - The check is on the *resolved* output path, so it covers the default - derived name (``mols_AIMNET_E.sdf``) exactly as it covers an explicit - ``-o``. A second ``auto3d energy mols.sdf`` therefore stops rather than - silently replacing the first run's results. - - ``os.path.exists`` follows symlinks, which is the behavior wanted here: a - symlink pointing at a real file is a file the write would destroy. A - dangling symlink reports False and is overwritten, matching what the - writer would do anyway. - - Args: - out_path: The resolved path the caller is about to write, or None - when the caller has nothing to write. - overwrite: True to allow clobbering an existing file (``--force`` on - the CLI, ``overwrite=True`` in the Python API). - - Raises: - ConfigurationError: `out_path` exists and `overwrite` is False. - """ - if out_path is None or overwrite: - return - - if os.path.exists(out_path): - raise ConfigurationError( - f"{out_path} already exists. Pass --force/-f to overwrite, or " - "choose a different -o path. (Python API: pass overwrite=True.)", - # No hint: the message above already states both ways out, and - # ConfigurationError's class hint ("run auto3d config init") has - # nothing to do with an -o collision. "" suppresses it; None - # would have meant "unset" and let the class hint through. - hint="", - ) - - def check_input(args: Any) -> None: """Check the input file and give recommendations. diff --git a/src/Auto3D/workflow.py b/src/Auto3D/workflow.py index a0259761..b49ee1c9 100644 --- a/src/Auto3D/workflow.py +++ b/src/Auto3D/workflow.py @@ -15,17 +15,13 @@ from Auto3D.chunk_manager import ChunkManager from Auto3D.config import Auto3DOptions, optimizer_worker_indices from Auto3D.exceptions import ConfigurationError, FileFormatError, OptimizationError +from Auto3D.id_mapping import decode_ids, encode_ids from Auto3D.model_factory import ModelFactory from Auto3D.models.preflight import preflight_model from Auto3D.torch_config import TorchConfig, configure_torch -from Auto3D.utils.file_ops import ( - decode_ids, - encode_ids, - find_ids_not_in_sdf, - find_smiles_not_in_sdf, - reorder_sdf, -) from Auto3D.utils.logging_config import get_logger +from Auto3D.utils.reconciliation import find_ids_not_in_sdf, find_smiles_not_in_sdf +from Auto3D.utils.sdf_io import reorder_sdf from Auto3D.utils.validation import check_input, check_valid_configuration from Auto3D.workflow_workers import ( isomer_wrapper, @@ -458,6 +454,20 @@ def _check_exit(self, proc: mp.Process, label: str) -> None: label, proc.exitcode, ) + def _log_both(self, msg: str, *, warning: bool = False) -> None: + """Emit ``msg`` to both the module logger and this run's Auto3D.log. + + ``self.logger`` (the per-run file logger ``_setup_logging`` attaches) + is ``None`` until logging starts, so every call site used to guard the + per-run copy with its own ``if self.logger:`` right next to an + identical module-level call -- four times (``_finalize_output``, + ``_reconcile_output``, and twice in ``_log_timing``), three at + ``info`` and one at ``warning``. + """ + (logger.warning if warning else logger.info)(msg) + if self.logger: + (self.logger.warning if warning else self.logger.info)(msg) + def _supervise_with_progress( self, p1: mp.Process, p2s: list[mp.Process], progress_queue: Queue[dict] ) -> None: @@ -511,9 +521,11 @@ def _finalize_output(self, start_time: float) -> str: """ # Combine all job outputs using pathlib glob output_files = list(self.job_dir.glob("job*/*_3d.sdf")) + # Computed once, up front: both failure messages below name it, and + # at most one of the two `raise`s below it can ever execute. + log_path = self.job_dir / "Auto3D.log" if not output_files: - log_path = self.job_dir / "Auto3D.log" raise OptimizationError( "No chunk produced a 3D structure output file, so no 3D " "structure converged. The model was already verified " @@ -533,7 +545,6 @@ def _finalize_output(self, start_time: float) -> str: combined_data.extend(file_path.read_text().splitlines(keepends=True)) if not any(line.strip() == "$$$$" for line in combined_data): - log_path = self.job_dir / "Auto3D.log" raise OptimizationError( "No 3D structure converged. Every chunk produced an output " "file, but none of them contain a converged structure. The " @@ -561,9 +572,7 @@ def _finalize_output(self, start_time: float) -> str: # Cleanup temporary files (input_path is unlinked in run()'s finally) path_combined.unlink() - logger.info(f"Output path: {path_output}") - if self.logger: - self.logger.info(f"Output path: {path_output}") + self._log_both(f"Output path: {path_output}") # Reconcile inputs against outputs (C7): a molecule that vanished # mid-pipeline must leave a trace. Compare the ORIGINAL input @@ -611,9 +620,7 @@ def _reconcile_output(self, path_output: str) -> None: f"{len(self.failures)} input molecule(s) produced no output " f"and were not reported anywhere else: {sorted(self.failures)}" ) - logger.warning(msg) - if self.logger: - self.logger.warning(msg) + self._log_both(msg, warning=True) def _log_timing(self, start_time: float) -> None: """Log pipeline execution time. @@ -621,9 +628,7 @@ def _log_timing(self, start_time: float) -> None: Args: start_time: Pipeline start time. """ - logger.info("Energy unit: Hartree if implicit.") - if self.logger: - self.logger.info("Energy unit: Hartree if implicit.") + self._log_both("Energy unit: Hartree if implicit.") elapsed_minutes = int((time.time() - start_time) / 60) @@ -634,6 +639,4 @@ def _log_timing(self, start_time: float) -> None: remaining = elapsed_minutes - hours * 60 msg = f"Program running time: {hours} hour(s) and {remaining} minute(s)" - logger.info(msg) - if self.logger: - self.logger.info(msg) + self._log_both(msg) diff --git a/src/Auto3D/workflow_workers.py b/src/Auto3D/workflow_workers.py index cd3b9a45..f3b348ab 100644 --- a/src/Auto3D/workflow_workers.py +++ b/src/Auto3D/workflow_workers.py @@ -24,10 +24,10 @@ from Auto3D.batch_opt.batchopt import optimizing from Auto3D.config import optimizer_worker_indices from Auto3D.isomers import IsomerEngineFactory +from Auto3D.job_layout import create_chunk_meta_names, housekeeping from Auto3D.model_factory import create_model from Auto3D.processors import TautomerProcessor from Auto3D.ranking import ranking -from Auto3D.utils.file_ops import create_chunk_meta_names, housekeeping if TYPE_CHECKING: from logging import LogRecord @@ -42,7 +42,7 @@ # Several call sites instead log through logging.getLogger("auto3d") # directly -- lowercase -- to work around the fact that "Auto3D.*" is a # different, case-distinct tree with no ancestor relationship to "auto3d" -# (Auto3D.workflow's self.logger, Auto3D.utils.chemistry's module logger, the +# (Auto3D.workflow's self.logger, Auto3D.clash_relief's module logger, the # stereochemistry-change warning in Auto3D.batch_opt.batchopt). Attaching a # QueueHandler onto BOTH trees here -- writing to the very same queue -- lets # get_logger(__name__) warnings reach the run log too, without touching any diff --git a/tests/helpers_pipeline_output.py b/tests/helpers_pipeline_output.py index 6993f7ce..1478ba1a 100644 --- a/tests/helpers_pipeline_output.py +++ b/tests/helpers_pipeline_output.py @@ -12,7 +12,7 @@ Why bounds and invariants, not expected numbers ----------------------------------------------- These checks are derived from the code paths they guard -- ``ranking``, -``batch_opt.batchopt``, ``ASE.geometry``, ``utils.file_ops`` and +``batch_opt.batchopt``, ``ASE.geometry``, ``utils.sdf_io`` and ``utils.energy`` -- not from observed output. Pinning an NNP's numerics to several decimals would make the tier fail on a model-version bump or a different BLAS, and a slow tier that fails on correct code is one people learn @@ -124,7 +124,7 @@ def formulas_from_smi_file(path: str | Path) -> dict[str, str]: """Map ``{molecule id: formula}`` for a whitespace-delimited .smi file. Deliberately a plain parse rather than a call to - ``Auto3D.utils.file_ops.iter_smi_records``: this is the *expectation* side + ``Auto3D.utils.smi_io.iter_smi_records``: this is the *expectation* side of the comparison, so reusing the production reader would let a bug in that reader cancel itself out. @@ -170,7 +170,7 @@ def base_molecule_id(name: str) -> str: ``decode_ids`` restored the user-facing id, so the only decoration that can remain is a ``@tautN`` tautomer suffix. Stripped exactly the way the pipeline's own reconciliation does it (``find_smiles_not_in_sdf`` / - ``find_ids_not_in_sdf`` in ``utils/file_ops.py``) so that the accounting + ``find_ids_not_in_sdf`` in ``utils/reconciliation.py``) so that the accounting assertion compares like with like. """ return name.split("@taut")[0].strip() diff --git a/tests/helpers_sync_count.py b/tests/helpers_sync_count.py new file mode 100644 index 00000000..f15f4285 --- /dev/null +++ b/tests/helpers_sync_count.py @@ -0,0 +1,223 @@ +# tests/helpers_sync_count.py +"""Count host<->device synchronization points without a GPU. + +Why this exists +--------------- +Auto3D's optimization loop runs up to 2000 steps per bucket, so a handful of +host-device serialization points per step is the difference between a +launch-bound and a compute-bound loop. CI has no GPU, so it can never *time* +that. It can, however, *count* it exactly -- which is what this module does, +and what ``test_optimization_engine_indexing.py`` asserts. + +On CUDA the sync-forcing operations reachable from this code are exactly four, +and all four are observable on CPU through ``TorchDispatchMode`` because the +sync is a property of the *operator*, not of the device: + +1. Boolean-mask advanced read ``x[bool_mask]`` dispatches ``aten.index.Tensor`` + with a bool index. ATen expands the mask via ``nonzero()`` and copies the + resulting element count to the host to size the output. Sync. +2. Boolean-mask advanced write ``x[bool_mask] = v`` dispatches + ``aten.index_put_`` with a bool index -- same ``nonzero()``. **Exception:** + ATen's ``canDispatchToMaskedFill`` fast path lowers it to ``masked_fill_`` + when the value is a CPU scalar with ``numel() == 1``, and that does *not* + sync. This module models that exception, because otherwise + ``oscillating_count[mask] = 0`` would be miscounted as a sync. +3. Scalar readback ``.item()`` / ``bool()`` / ``int()`` / ``float()`` dispatches + ``aten._local_scalar_dense``. Sync. +4. Device-to-host copy ``.cpu()`` / ``.to('cpu')`` / ``.tolist()`` / ``.numpy()`` + dispatches ``aten._to_copy`` across devices. Sync. + +Deliberately *not* syncs, and the basis of the fix these tests lock in: +``index_select``, ``index_copy_``, ``index_add_``, ``scatter_add_``, +``masked_fill``, ``torch.where``, ``Tensor.split`` with host-known sizes, and +integer-index advanced indexing ``x[int64_idx]``. ``x.shape[0]`` is host-side +metadata and is free. + +Honest limits of the method +--------------------------- +* Counting on CPU cannot price a sync; only a GPU can. See + ``benchmarks/bench_optimization_perf.py``. +* Rule 4 is unobservable in a CPU-only run (there is no second device), so + ``.cpu()`` on an already-CPU tensor shows up only through the + ``_local_scalar_dense`` of the subsequent ``int()``. The counts produced here + are therefore a *lower* bound on rule 4 and exact for rules 1-3, which is the + right direction: a regression can only ever be under-reported, never invented. +""" +from __future__ import annotations + +import collections +import traceback + +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +#: Label for a boolean-mask advanced read (``x[bool_mask]``). +BOOL_READ = "bool-mask READ (index.Tensor -> nonzero)" +#: Label for a boolean-mask advanced write (``x[bool_mask] = v``). +BOOL_WRITE = "bool-mask WRITE (index_put_ -> nonzero)" +#: Label for an explicit ``torch.nonzero`` call -- the *intended* one sync. +NONZERO = "explicit nonzero()" +#: Label for a scalar readback (``.item()``/``int()``/``bool()``/``float()``). +SCALAR_READBACK = "scalar readback (.item()/bool()/int())" +#: Label for a device-to-host copy. +D2H_COPY = "device-to-host copy" + +#: Every label that denotes a boolean-mask indexing op. These must be zero in +#: the hot loop after the M6 rewrite. +BOOL_MASK_LABELS = (BOOL_READ, BOOL_WRITE) + + +def _is_bool(t: object) -> bool: + return isinstance(t, torch.Tensor) and t.dtype in (torch.bool, torch.uint8) + + +def classify(func: object, args: tuple, kwargs: dict) -> str | None: + """Return a sync-kind label for ``func``, or ``None`` if it does not sync. + + Args: + func: The ``OpOverload`` handed to ``__torch_dispatch__``. + args: Positional args as dispatched (already flattened by ATen). + kwargs: Keyword args as dispatched. + + Returns: + One of the module-level labels, or ``None`` for a sync-free op. + """ + name = str(func) + + if name == "aten.index.Tensor": + idxs = args[1] if len(args) > 1 else () + if any(_is_bool(i) for i in idxs if i is not None): + return BOOL_READ + return None + + if name.startswith("aten.index_put_"): + idxs = args[1] if len(args) > 1 else () + if not any(_is_bool(i) for i in idxs if i is not None): + return None + val = args[2] if len(args) > 2 else None + # canDispatchToMaskedFill: a CPU scalar value with numel()==1 lowers to + # masked_fill_, which does not sync. Modelling this exception is what + # keeps `oscillating_count[mask] = 0` from being miscounted. + if isinstance(val, torch.Tensor) and val.numel() == 1 and val.device.type == "cpu": + return None + return BOOL_WRITE + + if name.startswith("aten.nonzero"): + return NONZERO + + if name == "aten._local_scalar_dense.default": + return SCALAR_READBACK + + if name.startswith("aten._to_copy"): + src = args[0] if args else None + dst_device = kwargs.get("device") + if (isinstance(src, torch.Tensor) and dst_device is not None + and torch.device(dst_device).type != src.device.type): + return D2H_COPY + return None + + return None + + +class SyncCounter(TorchDispatchMode): + """Dispatch-mode counter for host-device synchronization points. + + Use as a context manager around the code under test:: + + with SyncCounter() as counter: + n_steps(state, n=9, opttol=0.0, patience=10 ** 9) + assert counter.bool_mask_ops == 0 + + Attributes: + counts: ``Counter`` keyed by the module-level sync labels. + sites: ``Counter`` keyed by ``(label, "file:line source")``, populated + only when ``attribute=True`` (it walks the Python stack per op, so + it is slow and off by default). + """ + + def __init__(self, attribute: bool = False) -> None: + """Initialize the counter. + + Args: + attribute: Record the ``Auto3D`` source line responsible for each + sync. Useful for diagnosing a regression; costs a stack walk + per dispatched op, so leave it off for bulk counting. + """ + super().__init__() + self.counts: collections.Counter = collections.Counter() + self.sites: collections.Counter = collections.Counter() + self.attribute = attribute + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): # noqa: D105 + kwargs = kwargs or {} + kind = classify(func, args, kwargs) + if kind: + self.counts[kind] += 1 + if self.attribute: + self.sites[(kind, self._blame())] += 1 + return func(*args, **kwargs) + + @staticmethod + def _blame() -> str: + """Return ``file:line source`` of the innermost Auto3D frame.""" + for frame in reversed(traceback.extract_stack()): + if "/Auto3D/" in frame.filename and "helpers_sync_count" not in frame.filename: + where = frame.filename.split("/Auto3D/")[-1] + return f"{where}:{frame.lineno} {(frame.line or '').strip()}" + return "" + + @property + def total(self) -> int: + """Total sync-forcing ops counted.""" + return sum(self.counts.values()) + + @property + def bool_mask_ops(self) -> int: + """Number of boolean-mask read/write ops (the M6 regression signal).""" + return sum(self.counts[label] for label in BOOL_MASK_LABELS) + + def report(self) -> str: + """Human-readable breakdown, for use in assertion messages.""" + if not self.counts: + return "no sync-forcing ops" + lines = [f" {n:4d} {label}" for label, n in sorted(self.counts.items())] + if self.sites: + lines.append(" sites:") + lines += [f" {n:4d} {label}\n {site}" + for (label, site), n in sorted(self.sites.items(), key=lambda kv: -kv[1])] + return "\n".join(lines) + + +def count_graphs(fn, *args, dynamic: bool = True, fullgraph: bool = False) -> int: + """Compile ``fn`` with a counting backend and return the subgraph count. + + A ``torch.compile`` graph break inside a ``for`` loop makes Dynamo skip the + whole frame, which compiles to *zero* subgraphs rather than several -- so + "how many graphs" is the only reliable signal that the compile path works + at all. Controls: a clean function gives 1, one data-dependent branch in a + loop gives 0. + + Args: + fn: Callable (or ``nn.Module``) to compile. + *args: Arguments to invoke it with, once. + dynamic: Passed to ``torch.compile``. + fullgraph: Passed to ``torch.compile``; ``True`` turns any graph break + into an exception. + + Returns: + Number of subgraphs the backend was handed. + """ + import torch._dynamo as dynamo + + graphs: list = [] + + def backend(gm, example_inputs): + graphs.append(gm) + return gm.forward + + dynamo.reset() + try: + torch.compile(fn, backend=backend, dynamic=dynamic, fullgraph=fullgraph)(*args) + finally: + dynamo.reset() + return len(graphs) diff --git a/tests/test_ani2xt_atom_energies.py b/tests/test_ani2xt_atom_energies.py new file mode 100644 index 00000000..c1d7ab1f --- /dev/null +++ b/tests/test_ani2xt_atom_energies.py @@ -0,0 +1,376 @@ +# tests/test_ani2xt_atom_energies.py +"""Gate the M7 rewrite of ANI2xt's per-element energy loop. + +``ANI2xt.forward`` used to loop over its seven per-element networks doing +``if mask.any(): atom_energies[mask] = network(aev[mask])``. Measured with +``tests/helpers_sync_count.py``, that is **22 host-device synchronizations per +forward** with all seven elements present (7 for the guard, 7 masked reads, +7 masked writes, 1 for ``_validate_outputs``) and 16 with a drug-like four +elements -- on every step of every optimization. + +It was also uncompilable. ``if mask.any():`` is a data-dependent branch, and a +graph break *inside* a ``for`` loop gives Dynamo nowhere to place a resume +point, so it skipped the entire frame: ``compile_model=True`` produced **zero** +subgraphs for this model, not seven. Deleting the guard alone does not fix that, +because ``nonzero`` and boolean-mask indexing are dynamic-output-shape ops and +break the same way. Only a loop body with no data-dependent op at all compiles, +which is why the per-element indices are computed outside ``forward`` and passed +in. + +Every test here runs on CPU, needs no GPU, and -- because ``_atom_energies``, +``element_indices`` and ``self_atomic_energies`` are module-level functions +taking ``networks`` as a parameter rather than methods on a model that owns a +torchani AEV computer -- **needs no torchani**. That testability is the reason +for the extraction, not a side effect of it. + +What these tests cannot establish, and nothing here pretends otherwise: + +* Whether the *real* ``ANI2xt.forward``, with torchani's ``AEVComputer`` in the + frame, also reaches one subgraph. The AEV computer may break the graph on its + own, and any break inside the per-element loop re-triggers the whole-frame + skip. That needs torchani. +* Any wall-clock number. A sync costs only what it serializes, which depends on + the GPU and the batch size. See ``benchmarks/bench_optimization_perf.py``. +""" +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from Auto3D.batch_opt.ANI2xt_no_rep import ( + NUM_ELEMENTS, + _atom_energies, + element_indices, + self_atomic_energies, +) +from tests.helpers_sync_count import BOOL_MASK_LABELS, NONZERO, SyncCounter, count_graphs + +AEV_DIM = 8 + + +def _networks(num_elements: int = NUM_ELEMENTS, seed: int = 0) -> nn.ModuleList: + """Stand-in for ANI2xt's per-element MLPs. Shape-compatible, torchani-free.""" + torch.manual_seed(seed) + return nn.ModuleList([nn.Linear(AEV_DIM, 1) for _ in range(num_elements)]) + + +def _reference_atom_energies(networks, species_idx, aev): + """The pre-M7 loop: guard, boolean-mask read, boolean-mask write. + + Reimplemented here rather than imported, so bit-identity is measured against + an independent statement of the old behaviour. + """ + batch, atoms = species_idx.shape + out = torch.zeros(batch, atoms, device=aev.device, dtype=torch.float64) + for elem_idx, network in enumerate(networks): + mask = species_idx == elem_idx + if mask.any(): + out[mask] = network(aev[mask]).squeeze(-1).to(torch.float64) + return out + + +def _reference_element_indices(species_idx, num_elements=NUM_ELEMENTS): + """``nonzero`` per element -- what ``element_indices`` must reproduce exactly.""" + flat = species_idx.reshape(-1) + return [torch.nonzero(flat == elem, as_tuple=True)[0] for elem in range(num_elements)] + + +def _reference_self_energies(species_idx, shifts, num_elements=NUM_ELEMENTS): + """The inline self-energy loop that used to run on every forward.""" + out = torch.zeros(species_idx.shape[0], device=species_idx.device, dtype=torch.float64) + for elem_idx in range(num_elements): + counts = (species_idx == elem_idx).sum(dim=1).to(torch.float64) + out += counts * shifts[elem_idx] + return out + + +_SPECIES_CASES = { + "all seven elements": torch.tensor([[0, 1, 2, 3, 4, 5, 6, 0], [6, 5, 4, 3, 2, 1, 0, 1]]), + "only two of seven": torch.tensor([[0, 1, 0, 1, 1, 0, 0, 1], [1, 1, 0, 0, 0, 1, 1, 0]]), + "one element only": torch.full((2, 8), 3), + "with padding": torch.tensor([[0, 1, 2, -1, -1, -1, -1, -1], [3, 4, -1, -1, -1, -1, -1, -1]]), + "all padded": torch.full((2, 8), -1), + "single molecule single atom": torch.tensor([[5]]), + "out of range species": torch.tensor([[0, 6, 7, 99, -3, 2, 2, 2]]), + "empty batch": torch.zeros(0, 8, dtype=torch.long), +} + + +def _aev_for(species_idx: torch.Tensor, seed: int = 3) -> torch.Tensor: + torch.manual_seed(seed) + return torch.randn(*species_idx.shape, AEV_DIM) + + +class TestElementIndices: + """``element_indices`` reproduces seven ``nonzero`` calls with one readback.""" + + @pytest.mark.parametrize("label", list(_SPECIES_CASES)) + def test_matches_nonzero_reference(self, label): + """Same indices, same order, for every species pattern that can occur. + + Padded slots (``-1``) and out-of-range species get their own buckets and + are dropped, rather than being clamped into a neighbouring element's + network -- which is what would silently feed a stray atom to the wrong + MLP. + """ + species = _SPECIES_CASES[label].long() + expected = _reference_element_indices(species) + actual = element_indices(species) + + assert len(actual) == len(expected) == NUM_ELEMENTS + for elem_idx, (got, want) in enumerate(zip(actual, expected)): + assert got.dtype is torch.int64 + assert torch.equal(got, want), f"element {elem_idx} of {label}" + + def test_matches_nonzero_reference_over_random_patterns(self): + """A sweep, because the bucket arithmetic is easy to get subtly wrong. + + Species are drawn from ``[-3, NUM_ELEMENTS + 3)`` so padding, valid + indices and out-of-range values all occur, in every combination. + """ + for seed in range(200): + torch.manual_seed(seed) + batch = int(torch.randint(1, 6, (1,))) + atoms = int(torch.randint(1, 25, (1,))) + species = torch.randint(-3, NUM_ELEMENTS + 3, (batch, atoms)) + for got, want in zip(element_indices(species), + _reference_element_indices(species)): + assert torch.equal(got, want), f"seed {seed}" + + def test_performs_exactly_one_host_readback(self): + """One ``tolist``, no ``item``, no ``nonzero``. That is the whole point. + + Seven ``nonzero`` calls were seven synchronizations; this is one. The + readback is counted by wrapping ``Tensor.tolist``/``Tensor.item`` rather + than by dispatch mode, because a device-to-host copy is invisible in a + CPU-only process -- there is no second device for it to cross. + """ + calls = {"tolist": 0, "item": 0} + original_tolist, original_item = torch.Tensor.tolist, torch.Tensor.item + + def counting_tolist(self, *args, **kwargs): + calls["tolist"] += 1 + return original_tolist(self, *args, **kwargs) + + def counting_item(self, *args, **kwargs): + calls["item"] += 1 + return original_item(self, *args, **kwargs) + + species = torch.randint(-1, NUM_ELEMENTS, (16, 40)) + torch.Tensor.tolist, torch.Tensor.item = counting_tolist, counting_item + try: + element_indices(species) + finally: + torch.Tensor.tolist, torch.Tensor.item = original_tolist, original_item + + assert calls == {"tolist": 1, "item": 0} + + def test_does_no_nonzero_and_no_boolean_mask_indexing(self): + """The dispatch-level statement of the same claim.""" + counter = SyncCounter() + with counter: + element_indices(torch.randint(-1, NUM_ELEMENTS, (16, 40))) + assert counter.counts[NONZERO] == 0, counter.report() + assert counter.bool_mask_ops == 0, counter.report() + + def test_reference_really_does_seven_nonzeros(self): + """Prove the baseline, so the comparison above is not against nothing.""" + counter = SyncCounter() + with counter: + _reference_element_indices(torch.randint(-1, NUM_ELEMENTS, (16, 40))) + assert counter.counts[NONZERO] == NUM_ELEMENTS, counter.report() + + +class TestAtomEnergies: + """``_atom_energies`` is bit-identical to the masked loop it replaced.""" + + @pytest.mark.parametrize("label", list(_SPECIES_CASES)) + def test_matches_masked_reference(self, label): + """Exact equality, including the batch where only 2 of 7 elements appear. + + That case is the one that proves the deleted ``if mask.any():`` guard + protected nothing: for an absent element the index is empty, + ``network(empty)`` returns an empty tensor and ``index_copy`` with an + empty index is a no-op. + """ + species = _SPECIES_CASES[label].long() + aev = _aev_for(species) + networks = _networks() + batch, atoms = species.shape + + expected = _reference_atom_energies(networks, species, aev) + actual = _atom_energies( + networks, + aev.reshape(batch * atoms, AEV_DIM), + element_indices(species), + batch * atoms, + ).reshape(batch, atoms) + + assert actual.dtype is torch.float64 + assert torch.equal(actual, expected), label + + def test_padded_and_out_of_range_rows_stay_zero(self): + """Rows belonging to no element contribute exactly zero energy. + + Not "approximately zero": they are never written, so the initial zero + survives. A padded atom that picked up an energy would shift a molecule's + total by an amount depending on how much padding its bucket happened to + need. + """ + species = torch.tensor([[0, -1, 7, 3, -1]]) + aev = _aev_for(species) + result = _atom_energies( + _networks(), aev.reshape(5, AEV_DIM), element_indices(species), 5) + assert result[1] == 0.0 + assert result[2] == 0.0 + assert (result[[0, 3]] != 0.0).all() + + def test_does_no_boolean_mask_indexing(self): + """Zero masked reads, zero masked writes, zero ``nonzero`` in the loop.""" + species = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 0]]) + aev = _aev_for(species) + networks = _networks() + index = element_indices(species) + + counter = SyncCounter() + with counter: + _atom_energies(networks, aev.reshape(8, AEV_DIM), index, 8) + assert counter.total == 0, counter.report() + + def test_masked_reference_does_twenty_one_syncs(self): + """Pin the baseline this replaced: 7 guards + 7 reads + 7 writes. + + Without this the "22 -> 2 per forward" claim would rest on an unmeasured + memory of what the old code cost. + + Two atoms of every element, deliberately. With exactly *one* atom of an + element the assigned value has ``numel() == 1`` and ATen's + ``canDispatchToMaskedFill`` fast path lowers the write to + ``masked_fill_``, which does not sync -- so a sparser batch measures + fewer than 21 and the honest figure is "up to 21, and 21 for any + realistic molecule". + """ + species = torch.tensor([[0, 0, 1, 1, 2, 2, 3, 3], + [4, 4, 5, 5, 6, 6, 0, 0]]) + counter = SyncCounter() + with counter: + _reference_atom_energies(_networks(), species, _aev_for(species)) + assert counter.total == 3 * NUM_ELEMENTS, counter.report() + assert counter.bool_mask_ops == 2 * NUM_ELEMENTS, counter.report() + + def test_rewrite_does_one_sync_where_the_reference_did_twenty_one(self): + """End to end for the loop: 21 -> 1, and the 1 is the counts readback. + + The remaining synchronization is ``element_indices``' single host + readback, which a CPU-only process cannot observe as a device copy -- + hence the ``tolist`` count rather than a dispatch count. Add + ``_validate_outputs``' one scalar read, which is unchanged and lives in + the adapter, and the ANI2xt forward goes from 22 to 2. + """ + species = torch.tensor([[0, 0, 1, 1, 2, 2, 3, 3], + [4, 4, 5, 5, 6, 6, 0, 0]]) + aev = _aev_for(species) + networks = _networks() + + readbacks = 0 + original_tolist = torch.Tensor.tolist + + def counting_tolist(self, *args, **kwargs): + nonlocal readbacks + readbacks += 1 + return original_tolist(self, *args, **kwargs) + + counter = SyncCounter() + torch.Tensor.tolist = counting_tolist + try: + with counter: + index = element_indices(species) + _atom_energies(networks, aev.reshape(16, AEV_DIM), index, 16) + finally: + torch.Tensor.tolist = original_tolist + + assert counter.total == 0, counter.report() + assert readbacks == 1 + + +class TestCompilation: + """The graph-count claims, which are the load-bearing half of M7.""" + + def test_atom_energies_compiles_to_one_graph(self): + """One subgraph, and ``fullgraph=True`` succeeds. + + ``fullgraph=True`` is the sharper assertion: it turns any graph break + into an exception, so it cannot be satisfied by a frame that Dynamo + quietly skipped. + """ + species = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 0], [1, 1, 2, 2, 3, 3, 4, 4]]) + aev = _aev_for(species) + networks = _networks() + index = element_indices(species) + flat = aev.reshape(16, AEV_DIM) + + assert count_graphs(_atom_energies, networks, flat, index, 16) == 1 + assert count_graphs(_atom_energies, networks, flat, index, 16, + fullgraph=True) == 1 + + def test_masked_reference_compiles_to_zero_graphs(self): + """The control, without which "1 graph" means nothing. + + A data-dependent branch inside a ``for`` loop does not split the frame + into several graphs -- Dynamo abandons the frame, giving **zero**. This + is why the original report of "breaks into 7 graphs" was wrong in the + direction that made it worse, and why ``compile_model=True`` could not + have been speeding up this module. + """ + species = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 0], [1, 1, 2, 2, 3, 3, 4, 4]]) + aev = _aev_for(species) + networks = _networks() + + assert count_graphs(_reference_atom_energies, networks, species, aev) == 0 + + def test_masked_reference_cannot_satisfy_fullgraph(self): + """And it fails outright under ``fullgraph=True``, naming the reason.""" + species = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 0]]) + with pytest.raises(Exception, match="[Dd]ata-dependent"): + count_graphs(_reference_atom_energies, _networks(), species, + _aev_for(species), fullgraph=True) + + +class TestSelfAtomicEnergies: + """Hoisting a pure function of ``species`` out of the hot path.""" + + @pytest.mark.parametrize("label", list(_SPECIES_CASES)) + def test_matches_inline_reference(self, label): + """Bit-identical: the same seven terms summed in the same order.""" + species = _SPECIES_CASES[label].long() + shifts = torch.tensor( + [-0.5984, -38.0826, -54.7031, -75.1901, -99.8006, -398.1224, -460.1387], + dtype=torch.float64, + ) + assert torch.equal( + self_atomic_energies(species, shifts), + _reference_self_energies(species, shifts), + ) + + def test_is_independent_of_coordinates(self): + """Hoisting is only safe because nothing here depends on geometry. + + Stated as a test rather than a comment: the value is a function of + ``species`` and ``energy_shifts`` alone, so computing it once per bucket + instead of once per step cannot change any energy. + """ + species = torch.tensor([[0, 1, 2, 3, 4, 5, 6, -1]]) + shifts = torch.arange(NUM_ELEMENTS, dtype=torch.float64) * -1.5 + first = self_atomic_energies(species, shifts) + second = self_atomic_energies(species.clone(), shifts.clone()) + assert torch.equal(first, second) + assert first.dtype is torch.float64 + + def test_padded_slots_contribute_no_shift(self): + """``-1`` matches no element, so padding adds no self-energy.""" + shifts = torch.full((NUM_ELEMENTS,), -7.0, dtype=torch.float64) + unpadded = torch.tensor([[0, 1, 2]]) + padded = torch.tensor([[0, 1, 2, -1, -1]]) + assert torch.equal(self_atomic_energies(unpadded, shifts), + self_atomic_energies(padded, shifts)) diff --git a/tests/test_atomic_io.py b/tests/test_atomic_io.py new file mode 100644 index 00000000..f1bbe741 --- /dev/null +++ b/tests/test_atomic_io.py @@ -0,0 +1,261 @@ +"""The one shared in-place-rewrite helper, and the three call sites using it. + +Three functions in Auto3D rewrite a file they have just read -- ``reorder_sdf``, +``ASE.geometry._annotate_and_rewrite`` and +``utils.stereochemistry.amend_configuration_w`` -- and each grew its own +staging code. Two of them (the ``mkstemp`` pair) copied the target's permission +bits across; ``reorder_sdf`` did not, and it also used a predictable +``.reorder.tmp`` filename. So a 0600 SDF came back 0644 after a reorder: +a permission *loosening*, on exactly the path an ordinary +``auto3d run mols.smi`` takes. That is the divergence +:func:`Auto3D.utils.atomic_io.atomic_write_path` exists to remove. + +The table below is the contract, asserted once for the helper and once per call +site, so a site that stops using the helper fails here rather than silently +regressing to its own staging: + +* the temp file is a **sibling** of the target -- ``os.replace`` is only atomic + within one filesystem and raises ``EXDEV`` across them; +* the target's **mode is preserved** -- neither tightened (``mkstemp``'s 0600) + nor loosened (the process umask); +* **no temp file is left behind** when the body raises; +* the **original file is intact** when the body raises. +""" +from __future__ import annotations + +import os +import stat +from pathlib import Path + +import pytest +from rdkit import Chem +from rdkit.Chem import AllChem + +from Auto3D.utils.atomic_io import atomic_write_path + +# Captured before any test patches Chem.SDWriter (see tests/test_durability.py). +_real_sdwriter = Chem.SDWriter + + +def _write_sdf(path, names): + """Write one embedded ethanol conformer per name.""" + with _real_sdwriter(str(path)) as w: + for name in names: + mol = Chem.AddHs(Chem.MolFromSmiles("CCO")) + AllChem.EmbedMolecule(mol, randomSeed=42) + mol.SetProp("_Name", name) + mol.SetProp("ID", name) + w.write(mol) + + +def _mode(path) -> int: + return stat.S_IMODE(os.stat(str(path)).st_mode) + + +class TestAtomicWritePathHelper: + """The contract, on the helper itself.""" + + def test_yields_a_sibling_path(self, tmp_path): + target = tmp_path / "out.sdf" + target.write_text("original\n") + + with atomic_write_path(str(target), suffix=".sdf") as tmp: + assert Path(tmp).parent == Path(os.path.realpath(str(tmp_path))), ( + f"temp file {tmp} is not beside its target {target}; " + "os.replace would raise EXDEV whenever the two differ" + ) + Path(tmp).write_text("new\n") + + assert target.read_text() == "new\n" + + def test_parent_is_resolved_with_realpath_not_abspath(self, tmp_path): + """``abspath`` collapses ``..`` lexically and can pick the wrong mount. + + ``/scratch/link/../out.sdf`` where ``link`` points at another + filesystem: ``abspath`` says ``/scratch``, the replace destination is + somewhere else, and ``os.replace`` fails with EXDEV *after* a completed + run. Only the parent is resolved -- ``os.replace`` acts on the final + component itself, so following a symlinked target would pick the wrong + directory. + """ + real_dir = tmp_path / "real" + real_dir.mkdir() + link = tmp_path / "link" + link.symlink_to(real_dir, target_is_directory=True) + + target = link / ".." / "link" / "out.sdf" + Path(real_dir / "out.sdf").write_text("original\n") + + with atomic_write_path(str(target)) as tmp: + assert Path(tmp).parent == Path(os.path.realpath(str(real_dir))) + Path(tmp).write_text("new\n") + + assert (real_dir / "out.sdf").read_text() == "new\n" + + @pytest.mark.parametrize("mode", [0o600, 0o644, 0o640]) + def test_target_mode_is_preserved(self, tmp_path, mode): + """Neither tightened to mkstemp's 0600 nor loosened to the umask.""" + target = tmp_path / "out.sdf" + target.write_text("original\n") + os.chmod(target, mode) + + with atomic_write_path(str(target)) as tmp: + assert _mode(tmp) == mode, ( + "the staged temp file does not carry the target's mode, so " + "os.replace will change it" + ) + Path(tmp).write_text("new\n") + + assert _mode(target) == mode + + def test_missing_target_still_works(self, tmp_path): + """A target that does not exist yet has no mode to copy; do not fail.""" + target = tmp_path / "new.sdf" + + with atomic_write_path(str(target)) as tmp: + Path(tmp).write_text("new\n") + + assert target.read_text() == "new\n" + + def test_exception_leaves_no_temp_file_and_the_original_intact(self, tmp_path): + target = tmp_path / "out.sdf" + target.write_text("original\n") + + with pytest.raises(RuntimeError, match="disk full"): + with atomic_write_path(str(target)) as tmp: + Path(tmp).write_text("half a file") + raise RuntimeError("disk full") + + assert target.read_text() == "original\n", "the original was corrupted" + leftovers = sorted(p.name for p in tmp_path.iterdir() if p.name != "out.sdf") + assert not leftovers, f"temp files left behind: {leftovers}" + + def test_base_exception_also_cleans_up(self, tmp_path): + """KeyboardInterrupt mid-write must not leave a stray file either.""" + target = tmp_path / "out.sdf" + target.write_text("original\n") + + with pytest.raises(KeyboardInterrupt): + with atomic_write_path(str(target)) as tmp: + Path(tmp).write_text("half a file") + raise KeyboardInterrupt + + assert target.read_text() == "original\n" + leftovers = sorted(p.name for p in tmp_path.iterdir() if p.name != "out.sdf") + assert not leftovers, f"temp files left behind: {leftovers}" + + +class TestReorderSdfPreservesMode: + """``reorder_sdf`` is the call site the shared helper was written for. + + It staged through a predictable ``.reorder.tmp`` and never copied the + target's mode, so a 0600 input came back 0644 -- the process umask's idea of + a new file, applied to a file the user had deliberately restricted. This is + the one entry in the tripwire table that failed before the helper existed. + """ + + @pytest.mark.parametrize("mode", [0o600, 0o640, 0o644]) + def test_target_mode_survives_a_reorder(self, tmp_path, mode): + from Auto3D.utils.sdf_io import reorder_sdf + + sdf = tmp_path / "out.sdf" + smi = tmp_path / "in.smi" + _write_sdf(sdf, ["a", "b"]) + smi.write_text("CCO b\nCCO a\n") + os.chmod(sdf, mode) + + reorder_sdf(str(sdf), str(smi)) + + assert _mode(sdf) == mode, ( + f"reorder_sdf changed the output file's mode from {mode:04o} to " + f"{_mode(sdf):04o}" + ) + + def test_the_reordering_still_happens(self, tmp_path): + """Positive control: the mode test above must not pass vacuously.""" + from Auto3D.utils.sdf_io import reorder_sdf + + sdf = tmp_path / "out.sdf" + smi = tmp_path / "in.smi" + _write_sdf(sdf, ["a", "b"]) + smi.write_text("CCO b\nCCO a\n") + + mols = reorder_sdf(str(sdf), str(smi)) + assert [m.GetProp("_Name") for m in mols] == ["b", "a"] + on_disk = [ + m.GetProp("_Name") + for m in Chem.SDMolSupplier(str(sdf), removeHs=False) + if m is not None + ] + assert on_disk == ["b", "a"] + + def test_no_predictable_temp_name_is_used(self, tmp_path, monkeypatch): + """The staged name must come from ``mkstemp``, not from the target's. + + A predictable ``.reorder.tmp`` is guessable by any other process + sharing the directory, and two concurrent reorders of the same file + would stage onto each other's temp path. + """ + from Auto3D.utils.sdf_io import reorder_sdf + + sdf = tmp_path / "out.sdf" + smi = tmp_path / "in.smi" + _write_sdf(sdf, ["a"]) + smi.write_text("CCO a\n") + + seen: list[str] = [] + real = Chem.SDWriter + + def spy(path, *a, **k): + seen.append(str(path)) + return real(path, *a, **k) + + monkeypatch.setattr(Chem, "SDWriter", spy) + reorder_sdf(str(sdf), str(smi)) + + assert seen, "reorder_sdf opened no writer" + assert not any(p.endswith(".reorder.tmp") for p in seen), ( + f"reorder_sdf still stages through a predictable name: {seen}" + ) + + +class TestAmendConfigurationPreservesMode: + """The ``.smi`` call site keeps the property it already had.""" + + @pytest.mark.parametrize("mode", [0o600, 0o644]) + def test_target_mode_survives(self, tmp_path, mode): + from Auto3D.utils.stereochemistry import amend_configuration_w + + smi = tmp_path / "enum.smi" + smi.write_text("CC[C@H](N)O mol_a_0\n") + os.chmod(smi, mode) + + amend_configuration_w(str(smi)) + + assert _mode(smi) == mode + + +class TestAnnotateAndRewritePreservesMode: + """The ``opt_geometry`` rewrite pass keeps the property it already had.""" + + @pytest.mark.parametrize("mode", [0o600, 0o644]) + def test_target_mode_survives(self, tmp_path, mode): + from Auto3D.ASE.geometry import _annotate_and_rewrite + + out = tmp_path / "opt.sdf" + with _real_sdwriter(str(out)) as w: + mol = Chem.AddHs(Chem.MolFromSmiles("CCO")) + AllChem.EmbedMolecule(mol, randomSeed=42) + mol.SetProp("_Name", "a") + mol.SetProp("E_tot", "-1.5") + w.write(mol) + os.chmod(out, mode) + + _annotate_and_rewrite(str(out)) + + assert _mode(out) == mode + rewritten = [ + m for m in Chem.SDMolSupplier(str(out), removeHs=False) if m is not None + ] + assert len(rewritten) == 1 + assert rewritten[0].GetProp("E_tot(Hartree)") == "-1.5" diff --git a/tests/test_config.py b/tests/test_config.py index 776d0a7a..6c789132 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -408,7 +408,7 @@ class TestParallelEmbeddingIsReachable: Until 3.0.0 `use_parallel_embedding` was a constructor argument on the isomer engine with no route from `Auto3DOptions`, so no `main()` or `smiles2mols` run - could reach it and `isomers/parallel_embed.py` was reachable only from tests -- + could reach it and the parallel embedder was reachable only from tests -- which is why an audit listed that module as dead code. Asserting `Auto3DOptions(use_parallel_embedding=True).use_parallel_embedding diff --git a/tests/test_config_parity.py b/tests/test_config_parity.py index bb1e15dd..d50822e7 100644 --- a/tests/test_config_parity.py +++ b/tests/test_config_parity.py @@ -496,7 +496,7 @@ def run(self): f"{[m.GetProp('_Name') for m in mols]}" ) # M17: the second "CCO" is disambiguated to "_2" by - # smiles2smi's InChIKey-collision handling (utils/file_ops.py), and + # smiles2smi's InChIKey-collision handling (utils/smi_io.py), and # ranking.species_id must recover that suffix intact (rsplit on the # last two "_"-delimited components) rather than strip everything # after the FIRST underscore -- which would collapse both outputs diff --git a/tests/test_durability.py b/tests/test_durability.py index 07db026e..a8a3c818 100644 --- a/tests/test_durability.py +++ b/tests/test_durability.py @@ -10,69 +10,33 @@ and TestAmendConfigurationDurability must now also PASS -- Phase 6 gave both call sites the same tmp+os.replace staging. TestSameFileGuard was the tripwire for the other half of C14 (no out_path == path guard) and must now also PASS: -`Auto3D.utils.validation.check_output_not_input` refuses that case in all three -entry points, so the xfail(strict=True) marker it carried was removed with the -fix. Every class in this file is now a plain regression test -- this file -carries no xfail. +`Auto3D.utils.output_guard.check_output_not_input` refuses that case in all +three entry points, so the xfail(strict=True) marker it carried was removed +with the fix. Every class in this file is now a plain regression test -- this +file carries no xfail. + +All three call sites now stage through the one shared +`Auto3D.utils.atomic_io.atomic_write_path`, which is where the staging +*mechanism* -- sibling temp file, preserved permission bits, cleanup on any +exception -- is now asserted (`tests/test_atomic_io.py`). What stays here is the +end-to-end half: that an injected failure inside each real rewrite leaves the +real file intact. A `TestStagingLocation` class used to pin the sibling-and-mode +properties against `ASE.geometry._stage_beside`, which the shared helper +replaced. """ from __future__ import annotations from pathlib import Path +import os + import pytest import torch from rdkit import Chem from rdkit.Chem import AllChem -import os -import stat - -from Auto3D.utils.file_ops import reorder_sdf - - -class TestStagingLocation: - """The staged temp file must be a SIBLING of its target. - - `os.replace` is only atomic within one filesystem and raises - `OSError: [Errno 18] EXDEV` across them, so staging beside the target is - the property that makes the whole durability fix work. The end-to-end - durability tests below do NOT pin it: they patch `Chem.SDWriter` module- - globally, so the injected failure fires wherever the temp file happens to - live, and their leftover scans only read `job_dir`. Dropping `dir=` from - `_stage_beside` would leave every one of them green while breaking - `opt_geometry` on any box where the temp dir is a different mount from the - output directory -- a separate `/tmp` tmpfs being the common case. - """ - - def test_temp_file_is_created_beside_its_target(self, job_dir): - from Auto3D.ASE.geometry import _stage_beside +from Auto3D.utils.sdf_io import reorder_sdf - target = job_dir / "out.sdf" - target.write_text("placeholder\n") - - tmp_path = _stage_beside(str(target)) - try: - assert Path(tmp_path).parent == Path(os.path.realpath(str(job_dir))), ( - f"temp file {tmp_path} is not beside its target {target}; " - "os.replace would raise EXDEV whenever the two differ" - ) - finally: - os.unlink(tmp_path) - - def test_temp_file_inherits_the_target_mode(self, job_dir): - """mkstemp creates 0600 and os.replace carries the SOURCE mode, so - without this the rewrite would silently tighten every output file.""" - from Auto3D.ASE.geometry import _stage_beside - - target = job_dir / "out.sdf" - target.write_text("placeholder\n") - os.chmod(target, 0o644) - - tmp_path = _stage_beside(str(target)) - try: - assert stat.S_IMODE(os.stat(tmp_path).st_mode) == 0o644 - finally: - os.unlink(tmp_path) # Captured once, at import time, before any test monkeypatches Chem.SDWriter. # Constructing a real Chem.SDWriter on an existing path truncates it @@ -140,14 +104,22 @@ def __exit__(self, *a): assert sdf.read_bytes() == original, "the original SDF was corrupted" def test_no_temp_file_is_left_behind(self, job_dir, monkeypatch): - """A failed reorder must not leave a .tmp artifact next to the output. + """A failed reorder must leave no staged artifact next to the output. ``boom`` actually opens (and closes) the real writer at whatever path it is given before raising, so a genuine tmp file exists on disk ahead of the simulated crash -- otherwise there would be nothing for - the cleanup code (``file_ops.py``'s ``tmp_path.unlink()``) to ever + the cleanup code (``atomic_write_path``'s ``os.unlink``) to ever leave behind if that cleanup were removed, and this test would pass vacuously. + + The scan is by *exclusion* -- anything in the directory that is not one + of the two files this test created. It used to look for ``".tmp" in + name``, which stopped meaning anything once the staging moved to + ``atomic_write_path``: ``mkstemp`` names the file ``tmpXXXXXXXX.sdf``, + which contains no ``.tmp`` substring, so the old form would have gone + green whether or not the cleanup ran. Naming the expected files instead + cannot be defeated by a change of temp-file naming convention. """ sdf = job_dir / "out.sdf" smi = job_dir / "in.smi" @@ -163,7 +135,9 @@ def boom(path, *a, **k): with pytest.raises(Exception): reorder_sdf(str(sdf), str(smi)) - leftovers = [p.name for p in job_dir.iterdir() if ".tmp" in p.name] + leftovers = sorted( + p.name for p in job_dir.iterdir() if p.name not in {"out.sdf", "in.smi"} + ) assert not leftovers, f"temp files left behind: {leftovers}" @@ -1087,7 +1061,7 @@ def test_a_users_cwd_file_matching_the_sweep_globs_is_untouched( satisfied by a sweep that moved the file and happened to be interrupted before the deletion. """ - from Auto3D.utils.file_ops import housekeeping + from Auto3D.job_layout import housekeeping # The user's shell directory: `cd ~/project && auto3d run mols.smi`. project = job_dir / "project" diff --git a/tests/test_e_tot_units.py b/tests/test_e_tot_units.py index e3b591ba..9c8e915d 100644 --- a/tests/test_e_tot_units.py +++ b/tests/test_e_tot_units.py @@ -2,7 +2,7 @@ Before this, ``batch_opt.optimizing.run`` wrote ``E_tot`` in **eV** and ``ASE/geometry.opt_geometry`` rewrote the same tag in **Hartree**, while the -in-package consumers (``ranking``, ``filtering``, ``utils.chemistry``) all +in-package consumers (``ranking`` and both filters in ``filtering``) all hard-coded eV. The property name meant two different things depending on which entry point produced the file, and nothing in the file said which. diff --git a/tests/test_filter_unique.py b/tests/test_filter_unique.py new file mode 100644 index 00000000..e3ce2d35 --- /dev/null +++ b/tests/test_filter_unique.py @@ -0,0 +1,268 @@ +#!/usr/bin/env python +"""Tests for the legacy all-pairs ``Auto3D.filtering.filter_unique``. + +Kept as its own file because ``filter_unique`` is scheduled for removal once +``filter_unique_optimized`` is the single filter; deleting it then means +deleting this file, not surgery on a shared one. +""" +from __future__ import annotations + +import pytest # noqa: F401 +from rdkit import Chem +from rdkit.Chem import AllChem + +from Auto3D.filtering import filter_unique + + +class TestFilterUnique: + """Test the filter_unique function for RMSD-based duplicate filtering.""" + + def test_filter_identical_conformers(self): + """Test that identical conformers are filtered to one.""" + + mol = Chem.MolFromSmiles("CCO") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + mol.SetProp("Converged", "true") + + # Create identical copies + mol2 = Chem.Mol(mol) + mol2.SetProp("Converged", "true") + + mols = [mol, mol2] + unique_mols = filter_unique(mols, crit=0.3) + + # Should only keep one + assert len(unique_mols) == 1 + + def test_same_geometry_different_energy_kept(self): + """Identical geometry but distinct E_tot must be kept (energy guard). + + Heavy-atom RMSD ~= 0 but the two are distinct minima (the O-H rotamer + case); the energy guard must stop them collapsing into one. + """ + + mol = Chem.MolFromSmiles("CCO") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + mol.SetProp("Converged", "true") + mol.SetProp("E_tot", "-10.0") + + mol2 = Chem.Mol(mol) # identical geometry + mol2.SetProp("Converged", "true") + mol2.SetProp("E_tot", "-10.5") # |dE| >> tol + + unique_mols = filter_unique([mol, mol2], crit=0.3) + assert len(unique_mols) == 2 + + def test_missing_energy_falls_back_to_rmsd_only(self): + """Without E_tot the energy guard cannot apply -> RMSD-only dedup. + + Preserves the legacy behavior for callers that do not set E_tot. + """ + + mol = Chem.MolFromSmiles("CCO") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + mol.SetProp("Converged", "true") # no E_tot set + + mol2 = Chem.Mol(mol) + mol2.SetProp("Converged", "true") + + unique_mols = filter_unique([mol, mol2], crit=0.3) + assert len(unique_mols) == 1 + + def test_filter_different_conformers(self): + """Test that different conformers are kept.""" + + mol1 = Chem.MolFromSmiles("CCCCCC") # Hexane - flexible + mol1 = Chem.AddHs(mol1) + AllChem.EmbedMolecule(mol1, randomSeed=42) + mol1.SetProp("Converged", "true") + + mol2 = Chem.MolFromSmiles("CCCCCC") + mol2 = Chem.AddHs(mol2) + AllChem.EmbedMolecule(mol2, randomSeed=123) + mol2.SetProp("Converged", "true") + + # Generate very different conformers by using different seeds + # and moving atoms around + conf = mol2.GetConformer() + pos = conf.GetAtomPosition(0) + conf.SetAtomPosition(0, (pos.x + 0.5, pos.y, pos.z)) + + mols = [mol1, mol2] + unique_mols = filter_unique(mols, crit=0.3) + + # Should keep both (or at least not crash) + assert len(unique_mols) >= 1 + + def test_two_diastereomers_are_never_merged(self): + """A distinct compound must survive dedup, however close its geometry. + + The same guarantee ``tests/test_filtering.py`` asserts for + ``filter_unique_optimized``, asserted here because this is the other + duplicate filter and it applies the identical RMSD-plus-energy criterion. + Fixing one path and not the other would leave the defect reachable + through ``ConformerRanker(use_optimized_filtering=False)``. + + cis/trans-4-tert-butylcyclohexanol: heavy-atom RMSD between the two + diastereomers was measured at 0.300 A, i.e. at the 0.3 A default + threshold. ``crit`` is opened wide here so RMSD and energy both say + "duplicate" and the only thing that can keep the pair apart is the fact + that they are different compounds. + """ + from Auto3D.utils.energy import set_e_tot_from_ev + + def build(smiles: str) -> Chem.Mol: + mol = Chem.AddHs(Chem.MolFromSmiles(smiles)) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + mol.SetProp("Converged", "true") + set_e_tot_from_ev(mol, -10.0) # identical energies + return mol + + cis = build("O[C@H]1CC[C@@H](CC1)C(C)(C)C") + trans = build("O[C@H]1CC[C@H](CC1)C(C)(C)C") + assert Chem.MolToSmiles(cis) != Chem.MolToSmiles(trans), "test premise" + + assert len(filter_unique([cis, trans], crit=10.0)) == 2, ( + "the legacy filter merged two distinct diastereomers, so an input " + "molecule vanished from the output with no record" + ) + + def test_duplicate_conformers_of_one_stereoisomer_still_collapse(self): + """The other half: the stereo guard must narrow dedup, not disable it.""" + from Auto3D.utils.energy import set_e_tot_from_ev + + def build() -> Chem.Mol: + mol = Chem.AddHs(Chem.MolFromSmiles("O[C@H]1CC[C@@H](CC1)C(C)(C)C")) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + mol.SetProp("Converged", "true") + set_e_tot_from_ev(mol, -10.0) + return mol + + assert len(filter_unique([build(), build()], crit=10.0)) == 1, ( + "duplicate conformers of one stereoisomer survived, so the stereo " + "guard has switched dedup off rather than narrowing it" + ) + + def test_filter_unconverged_removed(self): + """Test that unconverged structures are removed.""" + + mol1 = Chem.MolFromSmiles("CCO") + mol1 = Chem.AddHs(mol1) + AllChem.EmbedMolecule(mol1, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol1) + mol1.SetProp("Converged", "true") + + mol2 = Chem.MolFromSmiles("CCO") + mol2 = Chem.AddHs(mol2) + AllChem.EmbedMolecule(mol2, randomSeed=123) + mol2.SetProp("Converged", "false") # Not converged + + mols = [mol1, mol2] + unique_mols = filter_unique(mols, crit=0.3) + + # Only converged one should remain + assert len(unique_mols) == 1 + assert unique_mols[0].GetProp("Converged").lower() == "true" + + def test_filter_empty_list(self): + """Test filtering empty list returns empty list.""" + + unique_mols = filter_unique([], crit=0.3) + assert len(unique_mols) == 0 + + def test_filter_custom_threshold(self): + """Test that custom RMSD threshold works.""" + + mol1 = Chem.MolFromSmiles("CCO") + mol1 = Chem.AddHs(mol1) + AllChem.EmbedMolecule(mol1, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol1) + mol1.SetProp("Converged", "true") + + mol2 = Chem.Mol(mol1) + mol2.SetProp("Converged", "true") + + mols = [mol1, mol2] + + # With very small threshold, might keep both + unique_mols_small = filter_unique(mols, crit=0.0001) + # With large threshold, definitely keep only one + unique_mols_large = filter_unique(mols, crit=10.0) + + # Large threshold should definitely merge identical mols + assert len(unique_mols_large) == 1 + # A tighter threshold can never merge MORE than a looser one -- the + # discarded half of this test's own computation, now actually checked. + assert len(unique_mols_small) >= len(unique_mols_large) + + def test_filter_unique_removehs_is_linear_and_nondestructive(self, monkeypatch): + """Legacy filter_unique strips Hs once per molecule (not per comparison) and + returns the originals with explicit H + exact positions intact.""" + import numpy as np + from rdkit import Chem + from rdkit.Chem import AllChem + + from Auto3D import filtering + + base = Chem.AddHs(Chem.MolFromSmiles("CCCCO")) + cids = AllChem.EmbedMultipleConfs(base, numConfs=5, randomSeed=1) + mols = [] + for cid in cids: + m = Chem.Mol(base, confId=int(cid)) + m.SetProp("Converged", "true") + mols.append(m) + n_atoms = base.GetNumAtoms() + orig_pos = {id(m): m.GetConformer().GetPositions().copy() for m in mols} + + calls = {"n": 0} + real_removehs = filtering.Chem.RemoveHs + + def counting(mol, *a, **k): + calls["n"] += 1 + return real_removehs(mol, *a, **k) + + monkeypatch.setattr(filtering.Chem, "RemoveHs", counting) + + result = filtering.filter_unique(mols, crit=0.01) + assert calls["n"] == len(mols) # once per input, never per pair + assert len(result) == len(mols) + for m in result: + assert m.GetNumAtoms() == n_atoms + assert any(a.GetAtomicNum() == 1 for a in m.GetAtoms()) + assert np.array_equal(m.GetConformer().GetPositions(), orig_pos[id(m)]) + + def test_rmsd_failure_keeps_both(self, monkeypatch): + """An incomparable pair (RMSD raises) must NOT be treated as a duplicate. + + When GetBestRMS raises RuntimeError, filter_unique must treat the pair + as distinct (rmsd = inf) and keep both, mirroring the fix already in + filtering._filter_within_cluster. The previous behavior (rmsd = 0) + made distinct conformers look like perfect duplicates and dropped one. + """ + from Auto3D import filtering + + def make(name): + m = Chem.AddHs(Chem.MolFromSmiles("CCO")) + AllChem.EmbedMolecule(m, randomSeed=abs(hash(name)) % 1000) + AllChem.MMFFOptimizeMolecule(m) + m.SetProp("_Name", name) + m.SetProp("Converged", "true") + return m + + def boom(*args, **kwargs): + raise RuntimeError("GetBestRMS failed") + + # filter_unique calls rdMolAlign.GetBestRMS via the filtering module. + monkeypatch.setattr(filtering.rdMolAlign, "GetBestRMS", boom) + + mols = [make("a"), make("b")] + unique_mols = filtering.filter_unique(mols, crit=0.3) + assert len(unique_mols) == 2 # incomparable pair must NOT be dropped diff --git a/tests/test_filtering.py b/tests/test_filtering.py index 3ea485eb..6974ca7c 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -333,7 +333,7 @@ def test_small_energy_window_creates_separate_clusters(self): class TestMissingEnergyPropertyMustNotCrash: """filter_unique_optimized must tolerate a record with no 'E_tot', the - way the legacy ``utils.chemistry.filter_unique`` already does. + way the legacy ``filtering.filter_unique`` already does. KNOWN DEFECT (found during cluster E brainstorming, not fixed by this lane): ``filtering.py:75`` sorts the valid-mols list by @@ -341,7 +341,7 @@ class TestMissingEnergyPropertyMustNotCrash: a molecule with no usable 'E_tot' property. ``_filter_within_cluster``'s own energy guard, two dozen lines later in the same module, instead uses the tolerant ``try_e_tot_ev`` and treats a missing energy as "fall back - to RMSD only". ``utils.chemistry.filter_unique`` (the OTHER conformer + to RMSD only". ``filtering.filter_unique`` (the OTHER conformer filter, sharing the same duplicate criterion since 4.0.1) also uses ``try_e_tot_ev`` throughout and does not crash on this input. So the two filters diverge on malformed input: the same list of mols that @@ -359,7 +359,7 @@ class TestMissingEnergyPropertyMustNotCrash: "filtering.py:75 sorts by e_tot_ev (raises KeyError for a mol " "with no 'E_tot' property) instead of the tolerant try_e_tot_ev " "that _filter_within_cluster's own energy guard and the legacy " - "utils.chemistry.filter_unique both use -- the two conformer " + "filtering.filter_unique both use -- the two conformer " "filters disagree on malformed input (cluster E brainstorm defect)." ), ) @@ -391,7 +391,7 @@ def test_matches_original_for_simple_case(self): real Auto3D contract). Use genuinely distinct conformers of one molecule so RMSD is well-defined and comparable across both implementations. """ - from Auto3D.utils.chemistry import filter_unique + from Auto3D.filtering import filter_unique def conformer(seed: float, energy_ev: float) -> Chem.Mol: m = Chem.AddHs(Chem.MolFromSmiles("CCCCCCO")) # flexible chain diff --git a/tests/test_fire_optimizer.py b/tests/test_fire_optimizer.py index 91306fea..563631d8 100644 --- a/tests/test_fire_optimizer.py +++ b/tests/test_fire_optimizer.py @@ -172,8 +172,8 @@ def test_fire_clean_subsets_state(self): optimizer(coord, forces) # Clean to keep only first 2 molecules - mask = torch.tensor([True, True, False, False]) - result = optimizer.clean(mask) + keep = torch.tensor([0, 1]) + result = optimizer.clean(keep) assert result is True assert optimizer.v.shape[0] == 2 @@ -195,36 +195,72 @@ def test_fire_clean_preserves_correct_molecules(self): original_dt1 = optimizer.dt[1].clone() # Keep only molecules 1 and 2 (indices 1, 2) - mask = torch.tensor([False, True, True]) - optimizer.clean(mask) + optimizer.clean(torch.tensor([1, 2])) # Molecule 1's state should now be at index 0 assert torch.allclose(optimizer.v[0], original_v1) assert torch.allclose(optimizer.dt[0], original_dt1) - def test_fire_clean_all_false_results_empty(self): - """FIRE.clean with all-false mask should result in empty tensors.""" + def test_fire_clean_empty_index_results_empty(self): + """FIRE.clean with an empty index should result in empty tensors. + + Reached when every remaining structure leaves the active set in the same + step; an empty ``index_select`` is a valid no-op-shaped result, not an + error. + """ coord = torch.randn(3, 5, 3) optimizer = FIRE(coord) - mask = torch.tensor([False, False, False]) - optimizer.clean(mask) + optimizer.clean(torch.zeros(0, dtype=torch.long)) assert optimizer.v.shape[0] == 0 assert optimizer.Nsteps.shape[0] == 0 - def test_fire_clean_all_true_unchanged(self): - """FIRE.clean with all-true mask should preserve all molecules.""" + def test_fire_clean_full_index_unchanged(self): + """FIRE.clean with every index should preserve all molecules.""" coord = torch.randn(3, 5, 3) optimizer = FIRE(coord) - mask = torch.tensor([True, True, True]) - optimizer.clean(mask) + optimizer.clean(torch.arange(3)) assert optimizer.v.shape[0] == 3 assert optimizer.Nsteps.shape[0] == 3 + def test_fire_clean_rejects_a_boolean_mask(self): + """A boolean mask must fail loudly, because it would not fail otherwise. + + ``clean`` used to take a boolean mask and now takes an int64 index. That + is the one genuinely dangerous kind of signature change: ``index_select`` + accepts ``tensor([True, True, False, False])`` and reinterprets it as + indices ``[1, 1, 0, 0]``, so an out-of-tree caller that was never updated + would get a silently permuted, wrong-length optimizer state instead of an + exception. Hence the explicit dtype check. + """ + optimizer = FIRE(torch.randn(4, 5, 3)) + with pytest.raises(ValueError, match="int64"): + optimizer.clean(torch.tensor([True, True, False, False])) + + def test_fire_clean_reindexes_in_the_order_given(self): + """State follows the index, position for position. + + ``n_steps`` always passes an ascending index from ``torch.nonzero``, but + pinning the general behaviour is what makes the bit-identity argument in + ``test_optimization_engine_indexing.py`` checkable: ``index_select`` + preserves order, so a converged molecule's removal never permutes the + survivors relative to the batch they are scattered back into. + """ + optimizer = FIRE(torch.randn(4, 5, 3)) + optimizer(torch.randn(4, 5, 3), torch.randn(4, 5, 3)) + before_dt = optimizer.dt.clone() + before_v = optimizer.v.clone() + + optimizer.clean(torch.tensor([2, 0])) + + assert torch.equal(optimizer.dt, torch.stack([before_dt[2], before_dt[0]])) + assert torch.equal(optimizer.v, torch.stack([before_v[2], before_v[0]])) + + class TestFIREMultipleSteps: """Tests for FIRE optimizer behavior over multiple steps.""" diff --git a/tests/test_id_mapping.py b/tests/test_id_mapping.py new file mode 100644 index 00000000..313fe297 --- /dev/null +++ b/tests/test_id_mapping.py @@ -0,0 +1,197 @@ +"""Tests for Auto3D.id_mapping module (encode_ids / decode_ids).""" +from pathlib import Path + +import pytest +from rdkit import Chem # noqa: F401 (several tests import it locally too) + +from Auto3D.id_mapping import decode_ids, encode_ids +from Auto3D.utils.sdf_io import count_sdf + +# Get the test files directory +TEST_DIR = Path(__file__).parent +FILES_DIR = TEST_DIR / "files" + + +class TestEncodeIdsSmiDenseIndex: + """encode_ids must use a dense, gap-free index for .smi inputs.""" + + def test_blank_lines_do_not_create_index_gaps(self, tmp_path): + smi = tmp_path / "in.smi" + # Blank lines interspersed: the old code used the file line number as the + # index, leaving gaps. The dense record counter must yield 0,1,2. + smi.write_text("CCO a\n\nCCC b\n\n\nCCCC c\n") + + new_path, mapping = encode_ids(str(smi)) + + assert sorted(mapping.values()) == [0, 1, 2] + ids = [line.split()[1] for line in Path(new_path).read_text().strip().split("\n")] + assert ids == ["0", "1", "2"] + + +class TestEncodeDecodeIds: + """Tests for encode_ids and decode_ids functions.""" + + def test_encode_smi_file(self, tmp_path): + """Test encoding IDs in a SMILES file.""" + input_file = tmp_path / "input.smi" + input_file.write_text("CCO mol_alpha\nCC mol_beta\nCCC mol_gamma\n") + + new_path, mapping = encode_ids(str(input_file)) + + assert mapping == {"mol_alpha": 0, "mol_beta": 1, "mol_gamma": 2} + assert Path(new_path).name == "input_encoded.smi" + + # Check encoded file content + content = Path(new_path).read_text() + assert "CCO 0" in content + assert "CC 1" in content + assert "CCC 2" in content + + def test_encode_sdf_file(self): + """Test encoding IDs in an SDF file.""" + sdf_path = str(FILES_DIR / "example.sdf") + + new_path, mapping = encode_ids(sdf_path) + + assert "mol1" in mapping + assert "mol2" in mapping + assert mapping["mol1"] == 0 + assert mapping["mol2"] == 1 + + # Clean up + Path(new_path).unlink(missing_ok=True) + + def test_encode_invalid_extension_raises(self, tmp_path): + """Test that invalid file extension raises ValueError.""" + input_file = tmp_path / "input.xyz" + input_file.write_text("invalid") + + with pytest.raises(ValueError, match="smi or sdf"): + encode_ids(str(input_file)) + + def test_encode_skips_blank_lines(self, tmp_path): + """Test that blank lines in SMILES file are skipped.""" + input_file = tmp_path / "input.smi" + input_file.write_text("CCO mol1\n\n \nCC mol2\n") + + new_path, mapping = encode_ids(str(input_file)) + + # mapping indices may not be sequential if blank lines are in between + assert len(mapping) == 2 + + # Clean up + Path(new_path).unlink(missing_ok=True) + + def test_encode_ids_rejects_duplicate_ids(self, tmp_path): + """Duplicate molecule IDs in a .smi file are rejected up front.""" + from Auto3D.exceptions import InputValidationError + + p = tmp_path / "dup.smi" + p.write_text("CCO mol1\nCCC mol1\n") + with pytest.raises(InputValidationError, match="[Dd]uplicate"): + encode_ids(str(p)) + + def test_encode_ids_rejects_missing_id(self, tmp_path): + """A .smi row without a whitespace-separated ID is rejected.""" + from Auto3D.exceptions import InputValidationError + + p = tmp_path / "noid.smi" + p.write_text("CCO\n") # no whitespace-separated ID + with pytest.raises(InputValidationError, match="ID"): + encode_ids(str(p)) + + def test_encode_ids_roundtrip_unique(self, tmp_path): + """Unique IDs encode cleanly and appear in the mapping.""" + p = tmp_path / "ok.smi" + p.write_text("CCO a\nCCC b\n") + _, mapping = encode_ids(str(p)) + assert set(mapping) == {"a", "b"} + + def test_encode_ids_rejects_blank_sdf_name(self, tmp_path): + """A molecule with a blank _Name in a .sdf file is rejected.""" + from rdkit.Chem import AllChem + + from Auto3D.exceptions import InputValidationError + + sdf = tmp_path / "blank.sdf" + with Chem.SDWriter(str(sdf)) as w: + m = Chem.AddHs(Chem.MolFromSmiles("CCO")) + AllChem.EmbedMolecule(m, randomSeed=1) + m.SetProp("_Name", "") # blank name + w.write(m) + with pytest.raises(InputValidationError): + encode_ids(str(sdf)) + + def test_encode_ids_refuses_to_overwrite_an_existing_file(self, tmp_path): + """The `_encoded.` name belongs to the user until proven + otherwise. + + The name is derived from the input, so `mols_encoded.smi` beside + `mols.smi` is an ordinary thing for a user to own -- and this function + used to open it for writing without a word. `WorkflowOrchestrator` + now redirects the encoded copy into its own job directory (see + `out_dir` below), but this check keeps the guarantee attached to the + function itself, so a caller taking the default location cannot + reintroduce the defect. + """ + from Auto3D.exceptions import ConfigurationError + + p = tmp_path / "mols.smi" + p.write_text("CCO a\n") + users_file = tmp_path / "mols_encoded.smi" + users_file.write_bytes(b"IRREPLACEABLE USER DATA\n") + + with pytest.raises(ConfigurationError, match="would overwrite"): + encode_ids(str(p)) + + assert users_file.read_bytes() == b"IRREPLACEABLE USER DATA\n" + + def test_encode_ids_writes_into_out_dir_when_given_one(self, tmp_path): + """`out_dir` moves the encoded copy somewhere the caller owns. + + This is how the run pipeline avoids the collision above entirely: it + passes the job directory it just created. The file name is unchanged, + only its directory -- downstream code (`_setup_job_directory`, + `decode_ids`) parses that name. + """ + p = tmp_path / "mols.smi" + p.write_text("CCO a\nCCC b\n") + staging = tmp_path / "staging" + staging.mkdir() + + new_path, mapping = encode_ids(str(p), out_dir=staging) + + assert Path(new_path).parent == staging + assert Path(new_path).name == "mols_encoded.smi" + assert mapping == {"a": 0, "b": 1} + assert not (tmp_path / "mols_encoded.smi").exists() + + +class TestNoneMolHardening: + """A None record yielded by SDMolSupplier must not crash decode_ids.""" + + def test_decode_ids_skips_none_records(self, tmp_path, monkeypatch): + """decode_ids must skip None records without raising.""" + + import Auto3D.id_mapping as id_mapping + + valid = Chem.MolFromSmiles("C") + valid.SetProp("_Name", "0") + valid.SetProp("ID", "0_conf1") + + monkeypatch.setattr( + id_mapping.Chem, "SDMolSupplier", lambda *a, **k: [valid, None] + ) + + # decode_ids expects a stem with at least two underscore parts. + sdf = tmp_path / "mols_3d_encoded.sdf" + sdf.write_text("placeholder") + + out = decode_ids(str(sdf), {"mol_a": 0}) + # Only the valid record is written; no AttributeError on the None. + written = count_sdf(out) + assert written == 1 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_import_boundaries.py b/tests/test_import_boundaries.py index c2f0f9f7..f07d5e37 100644 --- a/tests/test_import_boundaries.py +++ b/tests/test_import_boundaries.py @@ -398,7 +398,7 @@ def test_every_exported_name_is_documented_in_api_rst(): SRC_ROOT = pathlib.Path(__file__).resolve().parents[1] / "src" / "Auto3D" -# ``from Auto3D.utils import chemistry`` -- naming a submodule -- stays legal; +# ``from Auto3D.utils import energy`` -- naming a submodule -- stays legal; # it is a module reference, not a re-export. ``from Auto3D.utils import # hartree2ev`` does not. UTILS_SUBMODULES = frozenset( @@ -441,7 +441,8 @@ def test_no_src_module_imports_through_the_utils_barrel(): The barrel was never a coherent surface -- three of its eight submodules (``energy``, ``convergence``, ``stereo_check``) had no presence in it at all -- and ``check_connectivity`` was reached through the barrel in - ``filtering.py`` and through ``utils.chemistry`` in ``ranking.py``, two + ``filtering.py`` and through ``utils.chemistry`` in ``ranking.py`` (both now + name ``utils.connectivity``), two sibling modules disagreeing about the same function with nothing saying which was right. """ @@ -466,7 +467,7 @@ def test_utils_init_imports_nothing(): """``utils/__init__.py`` is docstring-only. Frozen deliberately, and not merely as the tail of the demolition: a later - cluster splits ``utils/file_ops.py`` and moves modules underneath this + cluster split ``utils/file_ops.py`` and moved modules underneath this package. Its own proposal was to "keep ``file_ops.py`` as a re-export shim so ``utils/__init__.py`` is untouched", which would rebuild the barrel one directory down. This test makes that fail here instead of being noticed @@ -514,18 +515,56 @@ def test_utils_package_exposes_no_names(): def test_utils_submodule_imports_still_work(): - """Emptying the barrel must not break ``from Auto3D.utils import chemistry``. + """Emptying the barrel must not break ``from Auto3D.utils import energy``. A submodule reference is not a re-export, several tests use this form, and ``__init__.py`` importing nothing is exactly the condition under which it is easy to assume otherwise. """ - from Auto3D.utils import chemistry, validation + from Auto3D.utils import energy, validation - assert chemistry.hartree2ev > 0 # a float constant, not a function + assert energy.hartree2ev > 0 # a float constant, not a function assert callable(validation.check_input) +def test_isomer_engine_does_not_import_the_isomers_package(): + """``Auto3D.isomers`` wraps ``isomer_engine``; the arrow may not point back. + + ``isomers.factory`` imports ``Auto3D.isomer_engine`` at module scope, so a + single import in the other direction closes a cycle. Until 4.0 that cycle + existed: the two adapter modules and ``factory.create_tautomer_engine`` + reached into ``isomer_engine``, and ``isomer_engine._run_parallel_embedding`` + reached back into ``isomers.parallel_embed``. It stayed latent only because + every edge was a function-scope import -- which is exactly the shape that + surfaces as an ``ImportError`` inside a ``spawn``ed worker and nowhere else, + since a spawned child re-imports from scratch in an order the parent never + exercised. + + Checked at **any** scope (``ast.walk``, not ``tree.body``) for that reason: + a function-scope import here would satisfy a module-scope-only check while + reintroducing precisely the latent cycle that was removed. ``parallel_embed`` + is now ``Auto3D.embedding``, which imports nothing from either side. + """ + path = SRC_ROOT / "isomer_engine.py" + tree = ast.parse(path.read_text(), filename=str(path)) + offenders = [] + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + module = _absolute_module(node, path) or "" + if module == "Auto3D.isomers" or module.startswith("Auto3D.isomers."): + offenders.append(f"line {node.lineno}: from {module} import ...") + elif isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "Auto3D.isomers" or alias.name.startswith( + "Auto3D.isomers." + ): + offenders.append(f"line {node.lineno}: import {alias.name}") + assert not offenders, ( + "isomer_engine.py imports from the Auto3D.isomers package it is wrapped " + "by, closing an import cycle:\n" + "\n".join(offenders) + ) + + def test_only_documented_subpackages_define_all(): """A subpackage may re-export only if api.rst documents it at that path.""" offenders = {} @@ -623,6 +662,54 @@ def test_importing_utils_validation_does_not_load_models(): ) +# Subprocess probe: what the split-out file-I/O modules cost to import. +_FILE_IO_PROBE_SOURCE = """ +import json, sys +import Auto3D.id_mapping # noqa: F401 +import Auto3D.job_layout # noqa: F401 +import Auto3D.utils.output_guard # noqa: F401 +import Auto3D.utils.reconciliation # noqa: F401 +import Auto3D.utils.sdf_io # noqa: F401 +import Auto3D.utils.smi_io # noqa: F401 +print(json.dumps({ + "torch": any(m == "torch" or m.startswith("torch.") for m in sys.modules), + "models": sorted( + m for m in sys.modules if m == "Auto3D.models" or m.startswith("Auto3D.models.") + ), +})) +""" + + +def test_file_io_modules_do_not_load_torch_or_the_model_tree(): + """Writing a ``.smi``/``.sdf`` file must not cost the neural-network stack. + + The six modules probed here are what ``utils/file_ops.py`` split into, and + the reason ``check_output_overwrite``/``check_output_not_input`` were lifted + out of ``utils/validation.py`` into the leaf ``utils/output_guard.py`` + first: the overwrite gate belongs on every one of these writers, and + reaching it through ``validation`` would have pulled in that module's + module-scope ``torch`` -- and, through the engine-name resolution it does + at function scope, the whole ``Auto3D.models`` tree -- for a caller that + only wanted to refuse clobbering a file. + + Subprocess, for the same reason as the probes above: ``conftest`` imports + every ``Auto3D`` submodule before any test runs, so in-process this would + pass unconditionally. + """ + proc = subprocess.run( + [sys.executable, "-c", _FILE_IO_PROBE_SOURCE], + capture_output=True, + text=True, + ) + assert proc.returncode == 0, f"probe failed:\n{proc.stdout}\n{proc.stderr}" + result = json.loads(proc.stdout.strip().splitlines()[-1]) + assert not result["torch"], "importing the .smi/.sdf writers pulled in torch" + assert not result["models"], ( + "importing the .smi/.sdf writers pulled in the models package: " + f"{result['models']}" + ) + + def test_validation_imports_torch_at_module_scope(): """The companion constraint, pinned so the leaf fix cannot overreach. diff --git a/tests/test_isomer_engine.py b/tests/test_isomer_engine.py index 393df1c1..773f077a 100644 --- a/tests/test_isomer_engine.py +++ b/tests/test_isomer_engine.py @@ -6,8 +6,7 @@ from rdkit import Chem from rdkit.Chem import rdMolAlign from Auto3D.isomer_engine import rd_isomer -from Auto3D.utils.file_ops import SDF2chunks -from Auto3D.utils.file_ops import count_sdf +from Auto3D.utils.sdf_io import SDF2chunks, count_sdf # Mark all tests in this module as slow (isomer embedding) pytestmark = pytest.mark.slow diff --git a/tests/test_isomer_engine_hardening.py b/tests/test_isomer_engine_hardening.py index ab900535..4d03191a 100644 --- a/tests/test_isomer_engine_hardening.py +++ b/tests/test_isomer_engine_hardening.py @@ -14,8 +14,9 @@ from rdkit import Chem from rdkit.Chem import AllChem +from Auto3D.constants import EV_TO_HARTREE from Auto3D.isomer_engine import RDKitIsomer, RDKitSdfIsomer, TautomerEngine -from Auto3D.utils.chemistry import calculate_conformer_count +from Auto3D.utils.molprops import calculate_conformer_count def _make_engine(tmp_path, smi_path, flipper=True, max_confs=None): @@ -315,8 +316,8 @@ def fake_pad(mols, adapter, device): monkeypatch.undo() written = list(Chem.SDMolSupplier(out, removeHs=False)) assert [m.GetProp("_Name") for m in written] == ["A", "B"] - assert float(written[0].GetProp("E_hartree")) == 10.0 * spe_mod.ev2hatree - assert float(written[1].GetProp("E_hartree")) == 20.0 * spe_mod.ev2hatree + assert float(written[0].GetProp("E_hartree")) == 10.0 * EV_TO_HARTREE + assert float(written[1].GetProp("E_hartree")) == 20.0 * EV_TO_HARTREE def test_calc_spe_all_filtered_does_not_crash(self, tmp_path, monkeypatch): """FIX B: an SDF whose only record is None must not raise the cryptic diff --git a/tests/test_isomers.py b/tests/test_isomers.py index 8a23fda1..703d6e75 100644 --- a/tests/test_isomers.py +++ b/tests/test_isomers.py @@ -1,18 +1,57 @@ -"""Unit tests for the isomers module.""" +"""Unit tests for the isomers package. + +The three adapter classes and the ``BaseIsomerEngine`` ABC these tests used to +exercise are gone: they existed only to copy ``IsomerEngineFactory.create``'s +keyword arguments into an attribute and then back out again into the real +engine, so their attribute assertions checked the copy rather than the mapping. +What replaced them is a single kwarg-mapping site inside ``create``, and the +tests below check it the way that actually pins the behavior -- by driving +``run()`` and asserting on the arguments the concrete engine +(``RDKitIsomer``/``RDKitSdfIsomer``/``oe_isomer``) is handed. +""" from __future__ import annotations import pytest -from Auto3D.isomers import ( - IsomerEngine, - IsomerEngineFactory, - TautomerEngine, - create_isomer_engine, - create_tautomer_engine, -) -from Auto3D.isomers.base import BaseIsomerEngine -from Auto3D.isomers.omega_adapter import OmegaIsomerAdapter -from Auto3D.isomers.rdkit_adapters import RDKitIsomerAdapter, RDKitSdfIsomerAdapter +from Auto3D.isomers import IsomerEngineFactory +from Auto3D.isomers.base import IsomerEngine, TautomerEngine +from Auto3D.isomers.factory import create_isomer_engine, create_tautomer_engine + + +@pytest.fixture +def spies(monkeypatch): + """Record the kwargs ``create(...).run()`` hands each concrete engine. + + Patched on ``Auto3D.isomers.factory``, the module that names them, so the + mapping under test is the one that runs. + """ + import Auto3D.isomers.factory as factory + + recorded: dict[str, dict] = {} + + class _FakeEngine: + def __init__(self, name, kwargs): + self._name = name + self._kwargs = kwargs + + def run(self): + recorded[self._name] = self._kwargs + return f"/{self._name}-output.sdf" + + def fake_rdkit(**kwargs): + return _FakeEngine("rdkit", kwargs) + + def fake_rdkit_sdf(**kwargs): + return _FakeEngine("rdkit_sdf", kwargs) + + def fake_oe_isomer(**kwargs): + recorded["omega"] = kwargs + return 0 + + monkeypatch.setattr(factory, "RDKitIsomer", fake_rdkit) + monkeypatch.setattr(factory, "RDKitSdfIsomer", fake_rdkit_sdf) + monkeypatch.setattr(factory, "oe_isomer", fake_oe_isomer) + return recorded class TestIsomerEngineProtocol: @@ -37,6 +76,15 @@ class NotAnEngine: obj = NotAnEngine() assert not isinstance(obj, IsomerEngine) + def test_what_create_returns_satisfies_the_protocol(self): + """The factory's return value is what callers type-check against.""" + engine = IsomerEngineFactory.create( + engine_type="rdkit_sdf", + input_path="/input.sdf", + output_path="/output.sdf", + ) + assert isinstance(engine, IsomerEngine) + class TestTautomerEngineProtocol: """Tests for TautomerEngine protocol.""" @@ -52,109 +100,158 @@ def run(self) -> None: assert isinstance(engine, TautomerEngine) -class TestBaseIsomerEngine: - """Tests for BaseIsomerEngine abstract class.""" - - def test_cannot_instantiate_directly(self): - """Test that BaseIsomerEngine cannot be instantiated directly.""" - with pytest.raises(TypeError): - BaseIsomerEngine( - input_path="/input.smi", - output_path="/output.sdf", - ) - - def test_subclass_stores_attributes(self): - """Test that subclass properly stores attributes.""" +class TestConstructionIsDeferredToRun: + """``create()`` must build nothing; ``run()`` builds and drives. - class ConcreteEngine(BaseIsomerEngine): - def run(self) -> str: - return self.output_path + Not a style point. ``RDKitIsomer.__init__`` calls ``self.rdk_tmp.mkdir()``, + so constructing the engine inside ``create()`` would move a filesystem side + effect -- and its ``FileExistsError`` on a second call with the same + ``job_dir`` -- from ``run()`` to ``create()``, where no caller expects it. + """ - engine = ConcreteEngine( + def test_create_builds_no_engine(self, spies): + IsomerEngineFactory.create( + engine_type="rdkit", input_path="/input.smi", output_path="/output.sdf", - max_confs=100, - threshold=0.5, - n_jobs=8, + job_dir="/job", ) + assert spies == {}, f"create() already built an engine: {sorted(spies)}" - assert engine.input_path == "/input.smi" - assert engine.output_path == "/output.sdf" - assert engine.max_confs == 100 - assert engine.threshold == 0.5 - assert engine.n_jobs == 8 + def test_create_does_not_touch_the_filesystem(self, tmp_path): + """The real ``RDKitIsomer``, unpatched: no ``rdk_tmp`` until ``run()``.""" + job_dir = tmp_path / "job" + job_dir.mkdir() - def test_default_values(self): - """Test default parameter values.""" + IsomerEngineFactory.create( + engine_type="rdkit", + input_path=str(tmp_path / "in.smi"), + output_path=str(tmp_path / "out.sdf"), + job_dir=str(job_dir), + ) - class ConcreteEngine(BaseIsomerEngine): - def run(self) -> str: - return self.output_path + assert not (job_dir / "rdk_tmp").exists(), ( + "create() created RDKitIsomer's working directory; construction " + "must stay deferred to run()" + ) - engine = ConcreteEngine( - input_path="/input.smi", + def test_run_returns_the_engines_output_path(self, spies): + engine = IsomerEngineFactory.create( + engine_type="rdkit_sdf", + input_path="/input.sdf", output_path="/output.sdf", ) + assert engine.run() == "/rdkit_sdf-output.sdf" - assert engine.max_confs is None - assert engine.threshold == 0.3 - assert engine.n_jobs == 4 +class TestCreateKwargMapping: + """Every argument ``create()`` accepts must reach the right engine argument.""" -class TestOmegaIsomerAdapter: - """Tests for OmegaIsomerAdapter class.""" - - def test_initialization(self): - """Test adapter initialization stores all parameters.""" - adapter = OmegaIsomerAdapter( - mode="classic", + def test_rdkit_mapping(self, spies): + IsomerEngineFactory.create( + engine_type="rdkit", input_path="/input.smi", + output_path="/output.sdf", smiles_enumerated="/enum.smi", smiles_reduced="/reduced.smi", smiles_hashed="/hashed.smi", - output_path="/output.sdf", + job_dir="/job", max_confs=50, threshold=0.25, + n_jobs=8, enumerate_isomers=False, - ) - - assert adapter.mode == "classic" - assert adapter.input_path == "/input.smi" - assert adapter.smiles_enumerated == "/enum.smi" - assert adapter.smiles_reduced == "/reduced.smi" - assert adapter.smiles_hashed == "/hashed.smi" - assert adapter.output_path == "/output.sdf" - assert adapter.max_confs == 50 - assert adapter.threshold == 0.25 - assert adapter.enumerate_isomers is False - - def test_default_values(self): - """Test adapter default parameter values.""" - adapter = OmegaIsomerAdapter( - mode="classic", + use_parallel_embedding=True, + parallel_embedding_threshold=5, + parallel_workers=2, + ).run() + + assert spies["rdkit"] == { + "smi": "/input.smi", + "smiles_enumerated": "/enum.smi", + "smiles_enumerated_reduced": "/reduced.smi", + "smiles_hashed": "/hashed.smi", + "enumerated_sdf": "/output.sdf", + "job_name": "/job", + "max_confs": 50, + "threshold": 0.25, + "np": 8, + "flipper": False, + "use_parallel_embedding": True, + "parallel_embedding_threshold": 5, + "parallel_workers": 2, + } + + def test_rdkit_defaults(self, spies): + IsomerEngineFactory.create( + engine_type="rdkit", input_path="/input.smi", - smiles_enumerated="/enum.smi", - smiles_reduced="/reduced.smi", - smiles_hashed="/hashed.smi", output_path="/output.sdf", - ) - - assert adapter.max_confs is None - assert adapter.threshold == 0.3 - assert adapter.enumerate_isomers is True - - def test_implements_protocol(self): - """Test that adapter implements IsomerEngine protocol.""" - adapter = OmegaIsomerAdapter( - mode="classic", + ).run() + + kwargs = spies["rdkit"] + assert kwargs["max_confs"] is None + assert kwargs["threshold"] == 0.3 + assert kwargs["np"] == 4 + assert kwargs["flipper"] is True + assert kwargs["use_parallel_embedding"] is False + assert kwargs["parallel_embedding_threshold"] == 10 + assert kwargs["parallel_workers"] == 4 + + def test_rdkit_sdf_mapping(self, spies): + IsomerEngineFactory.create( + engine_type="rdkit_sdf", + input_path="/input.sdf", + output_path="/output.sdf", + max_confs=7, + threshold=0.5, + n_jobs=3, + enumerate_isomers=False, + ).run() + + assert spies["rdkit_sdf"] == { + "sdf": "/input.sdf", + "enumerated_sdf": "/output.sdf", + "max_confs": 7, + "threshold": 0.5, + "np": 3, + "flipper": False, + } + + def test_omega_mapping(self, spies): + engine = IsomerEngineFactory.create( + engine_type="omega", input_path="/input.smi", + output_path="/output.sdf", smiles_enumerated="/enum.smi", smiles_reduced="/reduced.smi", smiles_hashed="/hashed.smi", - output_path="/output.sdf", + max_confs=50, + threshold=0.25, + enumerate_isomers=False, + mode="macrocycle", ) - - assert isinstance(adapter, IsomerEngine) + # oe_isomer is a function returning 0; the factory reports the path. + assert engine.run() == "/output.sdf" + + assert spies["omega"] == { + "mode": "macrocycle", + "input_f": "/input.smi", + "smiles_enumerated": "/enum.smi", + "smiles_reduced": "/reduced.smi", + "smiles_hashed": "/hashed.smi", + "output": "/output.sdf", + "max_confs": 50, + "threshold": 0.25, + "flipper": False, + } + + def test_omega_default_mode_is_classic(self, spies): + IsomerEngineFactory.create( + engine_type="omega", + input_path="/input.smi", + output_path="/output.sdf", + ).run() + assert spies["omega"]["mode"] == "classic" class TestCreateIsomerEngine: @@ -169,42 +266,43 @@ def test_unknown_engine_raises_error(self): output_path="/output.sdf", ) - def test_engine_type_case_insensitive(self): + def test_engine_type_case_insensitive(self, spies): """A *valid* engine name in an unexpected case must resolve to the - correct adapter, not merely fail to crash on an already-invalid name. + correct engine, not merely fail to crash on an already-invalid name. The previous version passed "UNKNOWN" -- invalid in any case -- so it could never have distinguished case normalization working from case normalization being entirely absent. """ for name in ("RDKit", "RDKIT", "rdkit"): - engine = create_isomer_engine( + spies.clear() + create_isomer_engine( name, input_path="/input.smi", output_path="/output.sdf", smiles_enumerated="/enum.smi", smiles_reduced="/reduced.smi", smiles_hashed="/hashed.smi", - ) - assert isinstance(engine, RDKitIsomerAdapter), name + ).run() + assert list(spies) == ["rdkit"], name - def test_omega_engine_creates_adapter(self): - """Test that 'omega' creates OmegaIsomerAdapter.""" - engine = create_isomer_engine( + def test_omega_engine_reaches_oe_isomer(self, spies): + """Test that 'omega' drives oe_isomer.""" + create_isomer_engine( "omega", input_path="/input.smi", output_path="/output.sdf", smiles_enumerated="/enum.smi", smiles_reduced="/reduced.smi", smiles_hashed="/hashed.smi", - ) + ).run() - assert isinstance(engine, OmegaIsomerAdapter) - assert engine.mode == "classic" + assert list(spies) == ["omega"] + assert spies["omega"]["mode"] == "classic" - def test_omega_engine_with_custom_mode(self): + def test_omega_engine_with_custom_mode(self, spies): """Test omega engine with custom mode.""" - engine = create_isomer_engine( + create_isomer_engine( "omega", input_path="/input.smi", output_path="/output.sdf", @@ -212,58 +310,46 @@ def test_omega_engine_with_custom_mode(self): smiles_reduced="/reduced.smi", smiles_hashed="/hashed.smi", mode="macrocycle", - ) + ).run() - assert engine.mode == "macrocycle" + assert spies["omega"]["mode"] == "macrocycle" class TestCreateIsomerEngineParallelEmbedding: """Tests for parallel embedding support in create_isomer_engine.""" - def test_rdkit_engine_parallel_embedding_default_off(self, tmp_path): + def test_rdkit_engine_parallel_embedding_default_off(self, spies): """Test that parallel embedding is off by default.""" - job_dir = tmp_path / "job" - job_dir.mkdir() - - engine = create_isomer_engine( + create_isomer_engine( "rdkit", input_path="/input.smi", output_path="/output.sdf", smiles_enumerated="/enum.smi", smiles_reduced="/reduced.smi", smiles_hashed="/hashed.smi", - job_dir=str(job_dir), - ) + job_dir="/job", + ).run() - # Now returns RDKitIsomerAdapter which wraps RDKitIsomer - from Auto3D.isomers.rdkit_adapters import RDKitIsomerAdapter - assert isinstance(engine, RDKitIsomerAdapter) - assert engine.use_parallel_embedding is False + assert spies["rdkit"]["use_parallel_embedding"] is False - def test_rdkit_engine_parallel_embedding_enabled(self, tmp_path): + def test_rdkit_engine_parallel_embedding_enabled(self, spies): """Test that parallel embedding can be enabled.""" - from Auto3D.isomers.rdkit_adapters import RDKitIsomerAdapter - - job_dir = tmp_path / "job" - job_dir.mkdir() - - engine = create_isomer_engine( + create_isomer_engine( "rdkit", input_path="/input.smi", output_path="/output.sdf", smiles_enumerated="/enum.smi", smiles_reduced="/reduced.smi", smiles_hashed="/hashed.smi", - job_dir=str(job_dir), + job_dir="/job", use_parallel_embedding=True, parallel_embedding_threshold=5, parallel_workers=2, - ) + ).run() - assert isinstance(engine, RDKitIsomerAdapter) - assert engine.use_parallel_embedding is True - assert engine.parallel_embedding_threshold == 5 - assert engine.parallel_workers == 2 + assert spies["rdkit"]["use_parallel_embedding"] is True + assert spies["rdkit"]["parallel_embedding_threshold"] == 5 + assert spies["rdkit"]["parallel_workers"] == 2 def test_rdkit_engine_parallel_embedding_enabled_actually_runs_parallel_path( self, tmp_path, monkeypatch @@ -276,7 +362,7 @@ def test_rdkit_engine_parallel_embedding_enabled_actually_runs_parallel_path( would be caught even though every attribute above still reports correctly. """ - import Auto3D.isomers.parallel_embed as parallel_embed_mod + import Auto3D.embedding as embedding_mod job_dir = tmp_path / "job" job_dir.mkdir() @@ -289,7 +375,7 @@ def spy(*args, **kwargs): calls["n"] += 1 return iter([]) # no conformers written; only the call matters - monkeypatch.setattr(parallel_embed_mod, "embed_conformers_parallel", spy) + monkeypatch.setattr(embedding_mod, "embed_conformers_parallel", spy) engine = create_isomer_engine( "rdkit", @@ -348,61 +434,17 @@ def test_available_engines(self): assert "rdkit_sdf" in engines assert "omega" in engines - def test_create_rdkit_engine(self, tmp_path): - """Test creating RDKit engine via factory.""" - job_dir = tmp_path / "job" - job_dir.mkdir() - - engine = IsomerEngineFactory.create( - engine_type="rdkit", - input_path="/input.smi", - output_path="/output.sdf", - smiles_enumerated="/enum.smi", - smiles_reduced="/reduced.smi", - smiles_hashed="/hashed.smi", - job_dir=str(job_dir), - ) - - assert isinstance(engine, RDKitIsomerAdapter) - assert isinstance(engine, BaseIsomerEngine) - - def test_create_rdkit_sdf_engine(self): - """Test creating RDKit SDF engine via factory.""" - engine = IsomerEngineFactory.create( - engine_type="rdkit_sdf", - input_path="/input.sdf", - output_path="/output.sdf", - ) - - assert isinstance(engine, RDKitSdfIsomerAdapter) - assert isinstance(engine, BaseIsomerEngine) - - def test_auto_select_rdkit_sdf_for_sdf_input(self): + def test_auto_select_rdkit_sdf_for_sdf_input(self, spies): """Test that rdkit auto-selects rdkit_sdf when input_format is sdf.""" - engine = IsomerEngineFactory.create( + IsomerEngineFactory.create( engine_type="rdkit", input_path="/input.sdf", output_path="/output.sdf", input_format="sdf", # This should trigger auto-selection - ) - - # Should get RDKitSdfIsomerAdapter, not RDKitIsomerAdapter - assert isinstance(engine, RDKitSdfIsomerAdapter) - - def test_create_omega_engine(self): - """Test creating Omega engine via factory.""" - engine = IsomerEngineFactory.create( - engine_type="omega", - input_path="/input.smi", - output_path="/output.sdf", - smiles_enumerated="/enum.smi", - smiles_reduced="/reduced.smi", - smiles_hashed="/hashed.smi", - mode="dense", - ) + ).run() - assert isinstance(engine, OmegaIsomerAdapter) - assert engine.mode == "dense" + # Should reach RDKitSdfIsomer, not RDKitIsomer + assert list(spies) == ["rdkit_sdf"] def test_unknown_engine_raises_error(self): """Test that unknown engine type raises ValueError.""" diff --git a/tests/test_job_layout.py b/tests/test_job_layout.py new file mode 100644 index 00000000..658a4f1f --- /dev/null +++ b/tests/test_job_layout.py @@ -0,0 +1,150 @@ +"""Tests for Auto3D.job_layout module.""" +from pathlib import Path + +import pytest # noqa: F401 (used by the __main__ guard below) + +from Auto3D.job_layout import create_chunk_meta_names, housekeeping + +# Get the test files directory +TEST_DIR = Path(__file__).parent +FILES_DIR = TEST_DIR / "files" + + +class TestHousekeeping: + """Tests for housekeeping function.""" + + def test_moves_files_except_output(self, tmp_path): + """Test that files are moved except for the output file.""" + job_dir = tmp_path / "job" + job_dir.mkdir() + + verbose_folder = tmp_path / "verbose" + verbose_folder.mkdir() + + # Create test files + (job_dir / "meta1.txt").write_text("meta1") + (job_dir / "meta2.txt").write_text("meta2") + output_file = job_dir / "output.sdf" + output_file.write_text("output") + + housekeeping(str(job_dir), str(verbose_folder), str(output_file)) + + # Output should still be in job_dir + assert output_file.exists() + # Meta files should be moved + assert (verbose_folder / "meta1.txt").exists() + assert (verbose_folder / "meta2.txt").exists() + + +class TestCreateChunkMetaNames: + """Tests for create_chunk_meta_names function.""" + + def test_generates_expected_paths(self): + """Test that all expected paths are generated.""" + result = create_chunk_meta_names("chunk1.smi", "/tmp/job") + + assert result["output"] == "/tmp/job/chunk1_3d.sdf" + assert result["optimized_og"] == "/tmp/job/chunk1_3d0.sdf" + assert result["output_taut"] == "/tmp/job/smi_taut.smi" + assert result["smiles_enumerated"] == "/tmp/job/smiles_enumerated.smi" + assert result["smiles_reduced"] == "/tmp/job/smiles_enumerated_reduced.smi" + assert result["smiles_hashed"] == "/tmp/job/smiles_enumerated_hashed.smi" + assert result["enumerated_sdf"] == "/tmp/job/smiles_enumerated.sdf" + assert result["sorted_sdf"] == "/tmp/job/enumerated_sorted.sdf" + assert result["housekeeping_folder"] == "/tmp/job/verbose" + assert result["path"] == "chunk1.smi" + assert result["dir"] == "/tmp/job" + + def test_handles_path_with_directory(self): + """Test that paths with directories work correctly.""" + result = create_chunk_meta_names("/data/input/chunk1.smi", "/output/job") + + assert result["output"] == "/output/job/chunk1_3d.sdf" + assert result["path"] == "/data/input/chunk1.smi" + + +class TestFileOpsIntegration: + """Integration tests for the job-layout helpers.""" + + def test_create_chunks_and_housekeeping_workflow(self, tmp_path): + """Test a typical workflow using both job_layout functions.""" + # Create job directory structure + job_dir = tmp_path / "job" + job_dir.mkdir() + + # Create meta names + meta = create_chunk_meta_names("input.smi", str(job_dir)) + + # Verify structure + assert "verbose" in meta["housekeeping_folder"] + + # Create verbose folder + Path(meta["housekeeping_folder"]).mkdir() + + # Create some intermediate files + Path(meta["smiles_enumerated"]).write_text("CCO mol1\n") + Path(meta["output"]).write_text("fake sdf output") + + # Run housekeeping - should move enumerated but not output + housekeeping( + str(job_dir), + meta["housekeeping_folder"], + meta["output"] + ) + + # Output should still exist + assert Path(meta["output"]).exists() + # Enumerated should be moved to verbose folder + assert (Path(meta["housekeeping_folder"]) / "smiles_enumerated.smi").exists() + + +def test_housekeeping_sweep_is_per_file_robust(tmp_path, monkeypatch): + """One unmovable file must not abandon the rest of the sweep. + + This guard used to live on a second loop that swept `oeomega_*` out of the + *process working directory*; that loop is gone (it destroyed user files -- + see `TestHousekeepingStaysInsideTheJobDirectory` in tests/test_durability.py) + and the OpenEye logfiles it collected now land inside the job directory, + where this loop picks them up. The robustness property moved with them: a + permission error, or a file that vanished under us, must leave a complete + `verbose` folder minus that one file rather than a half-populated one plus + a traceback out of `optim_rank_wrapper`'s blanket except. + """ + import os + + from Auto3D.job_layout import housekeeping + + job = tmp_path / "job" + job.mkdir() + dest = tmp_path / "verbose" + dest.mkdir() + + # Two logfiles in the job directory; the FIRST one encountered (by + # counter) will fail to move. + (job / "oeomega_a.log").write_text("a") + (job / "oeomega_b.log").write_text("b") + + real_move = __import__("shutil").move + call_count = {"n": 0} + + def flaky_move(src, dst): + call_count["n"] += 1 + if call_count["n"] == 1: + # Simulate the file having gone away underneath the sweep. + if os.path.exists(src): + os.remove(src) + raise OSError("already gone") + return real_move(src, dst) + + monkeypatch.setattr("Auto3D.job_layout.shutil.move", flaky_move) + + housekeeping(str(job), str(dest), str(job / "out.sdf")) # must not raise + + # Exactly one of the two logfiles must have been successfully moved: the + # sweep continued past the failure instead of stopping on it. + moved = list(dest.glob("oeomega_*.log")) + assert len(moved) == 1, f"Expected 1 moved file, got {[f.name for f in moved]}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_optimization_engine.py b/tests/test_optimization_engine.py index 71a7bce6..ecdb7d7e 100644 --- a/tests/test_optimization_engine.py +++ b/tests/test_optimization_engine.py @@ -275,9 +275,9 @@ class _ConstantForceNN: ``force_per_atom`` may be a scalar (every molecule feels the same force) or a ``{species: force}`` map for a heterogeneous batch. The map is keyed on species rather than on row index because ``n_steps`` gathers a SUBSET of - the batch once some molecules converge (``optimization_engine.py:168-177`` - gathers ``coord`` and ``numbers`` with the same ``not_converged`` mask), so - row 1 of a later step is not molecule 1. Keying on ``numbers`` -- gathered + the batch once some molecules converge (``_step_active_subset`` gathers + ``coord`` and ``numbers`` with the same ``active_idx``), so row 1 of a later + step is not molecule 1. Keying on ``numbers`` -- gathered alongside ``coord`` -- makes each force follow its own molecule. """ @@ -392,8 +392,8 @@ def test_convergence_outcome_never_depends_on_energy_stability(): # A heterogeneous batch, so the step loop actually performs a PARTIAL # subset gather. Every cell above uses one force for both molecules, so - # `not_converged` is always all-True or all-False and the gathers at - # optimization_engine.py:168-177 are no-ops. That leaves a reintroduction + # `not_converged` is always all-True or all-False and the gathers in + # `_step_active_subset` are no-ops. That leaves a reintroduction # bug invisible: a criterion whose per-molecule buffer is gathered with a # stale mask would let molecule 1, after molecule 0 converges and drops # out, read molecule 0's row and early-terminate at the wrong geometry. diff --git a/tests/test_optimization_engine_indexing.py b/tests/test_optimization_engine_indexing.py new file mode 100644 index 00000000..1bceaab6 --- /dev/null +++ b/tests/test_optimization_engine_indexing.py @@ -0,0 +1,834 @@ +# tests/test_optimization_engine_indexing.py +"""Gate the M6 host-device-sync reduction in ``n_steps``. + +``n_steps`` used to subset the batch with boolean masks: six masked reads, six +masked writes and four more inside ``FIRE.clean``, each of which forces a +GPU->CPU synchronization on CUDA because ATen has to ``nonzero()`` the mask and +copy the element count to the host to size the output. Measured with +``tests/helpers_sync_count.py``: **exactly 18 per step**, every step, up to 2000 +steps per bucket. + +The rewrite computes ``torch.nonzero(not_converged)`` **once** and feeds the +resulting int64 index to ``index_select`` / ``index_copy_``, which do not sync. +``nonzero`` is itself a sync -- that is the very mechanism by which boolean-mask +indexing syncs -- so this is **18 -> 2 per step**, not 18 -> 0: one ``nonzero`` +for the active subset and a second inside the step for ``FIRE.clean``, whose +mask is indexed within the active subset rather than the full batch. + +These tests assert two independent things: + +* **Results are unchanged.** ``TestEngineMatchesBooleanMaskReference`` runs a + test-local reimplementation of the *old* boolean-mask loop against production + ``n_steps`` in the same process and asserts ``torch.equal`` on every state + tensor. Same process and same hardware, so there is no cross-platform + float-determinism hazard -- which is exactly why a checked-in golden + trajectory file is *not* used here. +* **The win does not silently regress.** ``TestHotLoopDoesNotSync`` counts + dispatched ops and fails the moment someone reintroduces + ``state['coord'][mask]``. CI has no GPU and can never time this loop, but it + can count it exactly. +""" +from __future__ import annotations + +import pytest +import torch + +from Auto3D.batch_opt.fire_optimizer import FIRE +from Auto3D.batch_opt.optimization_engine import n_steps +from tests.helpers_sync_count import BOOL_MASK_LABELS, NONZERO, SyncCounter + +# --------------------------------------------------------------------------- # +# Hermetic potentials. No NNP, no model download, no GPU, CPU-only, exact. +# --------------------------------------------------------------------------- # + + +class _AnisotropicHarmonic: + """``E = sum_i k_i * r_i^2`` with a different ``k`` per molecule. + + The per-molecule force constant is what makes this a real test: molecules + converge at *different* steps, so the subset gather is a genuine partial + gather rather than a whole-batch no-op. A homogeneous batch hides exactly + the bug this file exists to catch (the same trap documented on + ``test_convergence_outcome_never_depends_on_energy_stability``). + + ``k`` is keyed on ``numbers[:, 0]`` so it follows a molecule through every + re-gather, which a positional lookup would not. + """ + + def __init__(self, k: dict[int, float]) -> None: + self.k = k + + def forward_batched(self, coord, numbers, charges, atom_mask=None): + kk = torch.tensor( + [self.k[int(numbers[row, 0])] for row in range(coord.shape[0])], + dtype=coord.dtype, + ).reshape(-1, 1, 1) + energy = (kk * coord ** 2).sum(dim=(1, 2)).to(torch.double) + forces = -2.0 * kk * coord + return energy, forces + + +class _ConstantForce: + """Force never decreases, so every molecule leaves via the oscillation drop. + + This is the only way to exercise the ``patience`` path, and (with a + ``patience`` larger than ``n``) the only way to keep every molecule active + for a known number of steps -- which the sync count below depends on. + """ + + def forward_batched(self, coord, numbers, charges, atom_mask=None): + forces = torch.zeros_like(coord) + forces[..., 0] = 0.5 + return torch.zeros(coord.shape[0], dtype=torch.double), forces + + +def _make_state(nn, batch: int, natoms: int, seed: int, start: float = 2.0) -> dict: + """Build a fresh ``n_steps`` state dict with reproducible coordinates.""" + torch.manual_seed(seed) + numbers = torch.arange(batch).reshape(batch, 1).expand(batch, natoms).contiguous() + return { + "coord": torch.rand(batch, natoms, 3, dtype=torch.float) * start, + "numbers": numbers.to(torch.long), + "charges": torch.zeros(batch, dtype=torch.long), + "nn": nn, + "converged_mask": torch.zeros(batch, dtype=torch.bool), + "fmax": torch.full((batch,), 999.0), + "energy": torch.full((batch,), float("inf"), dtype=torch.double), + } + + +# --------------------------------------------------------------------------- # +# T1 / T2 -- the primitive claims the rewrite rests on. +# --------------------------------------------------------------------------- # + +_SHAPES = { + "1d": (7,), + "2d-col": (7, 1), + "3d-coords": (7, 4, 3), +} +_MASKS = { + "all-true": [True] * 7, + "all-false": [False] * 7, + "single-true": [False, False, True, False, False, False, False], + "alternating": [True, False, True, False, True, False, True], + "random": [False, True, True, False, False, True, False], +} + + +def _sample(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + torch.manual_seed(11) + if dtype is torch.bool: + return torch.rand(shape) > 0.5 + if dtype is torch.long: + return torch.randint(-50, 50, shape) + return torch.randn(shape, dtype=dtype) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.long, torch.bool]) +@pytest.mark.parametrize("shape_name", list(_SHAPES)) +@pytest.mark.parametrize("mask_name", list(_MASKS)) +def test_index_select_matches_boolean_mask_read(dtype, shape_name, mask_name): + """``x.index_select(0, nonzero(m))`` equals ``x[m]``, bit for bit. + + Both emit rows in ascending index order (``nonzero`` returns sorted + indices), so the gathered subset is identical, not merely equivalent. This + is the read half of M6 and it must hold for every dtype in ``state``: + float coordinates, float64 energies, int64 counters, bool masks. + """ + x = _sample(_SHAPES[shape_name], dtype) + mask = torch.tensor(_MASKS[mask_name]) + idx = torch.nonzero(mask, as_tuple=True)[0] + + assert idx.dtype is torch.int64 + assert torch.equal(x.index_select(0, idx), x[mask]) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.long, torch.bool]) +@pytest.mark.parametrize("shape_name", list(_SHAPES)) +@pytest.mark.parametrize("mask_name", list(_MASKS)) +def test_index_copy_matches_boolean_mask_write(dtype, shape_name, mask_name): + """``dst.index_copy_(0, nonzero(m), v)`` equals ``dst[m] = v``, bit for bit. + + ``nonzero`` yields unique indices, so unlike ``index_add_`` / ``scatter_add_`` + there is no accumulation-order ambiguity to worry about. The all-false case + is included on purpose: an empty index is a valid no-op, which is what the + mid-loop "everything just converged" path relies on. + """ + shape = _SHAPES[shape_name] + mask = torch.tensor(_MASKS[mask_name]) + idx = torch.nonzero(mask, as_tuple=True)[0] + source = _sample((int(idx.numel()), *shape[1:]), dtype) + + expected = _sample(shape, dtype).clone() + expected[mask] = source + actual = _sample(shape, dtype).clone() + actual.index_copy_(0, idx, source) + + assert torch.equal(actual, expected) + + +def test_index_copy_and_index_put_dtype_strictness_measured(): + """Pin the exact dtype strictness of both write forms. Measured, not assumed. + + The intuition "``x[mask] = v`` silently casts, ``index_copy_`` raises" is only + half right, and the wrong half matters: + + * ``x[mask] = scalar`` **does** cast silently, because ATen's + ``canDispatchToMaskedFill`` fast path lowers a single-element value to + ``masked_fill_``. + * ``x[mask] = tensor`` with a mismatched dtype **already raises** -- + ``"Index put requires the source and destination dtypes match"`` -- in + both directions, narrowing and widening. + * ``index_copy_`` raises in every mismatched case, with its own message. + + So switching to ``index_copy_`` does not introduce a new class of failure for + tensor writes; it only removes the scalar fast path's silent cast. The + consequence for ``n_steps`` is the same either way: cast explicitly at every + write site, which is what it now does. + """ + mask = torch.tensor([True, True, False]) + index = torch.tensor([0, 1]) + + # Scalar value: index_put_ casts silently via masked_fill_. + scalar_dst = torch.zeros(3, dtype=torch.float32) + scalar_dst[torch.tensor([True, False, False])] = torch.ones(1, dtype=torch.float64) + assert scalar_dst.dtype is torch.float32 + assert scalar_dst[0] == 1.0 + + # Tensor value: BOTH forms raise, with different messages. + put_dst = torch.zeros(3, dtype=torch.float32) + with pytest.raises(RuntimeError, match="Index put requires the source and destination"): + put_dst[mask] = torch.ones(2, dtype=torch.float64) + + copy_dst = torch.zeros(3, dtype=torch.float32) + with pytest.raises(RuntimeError, match="expected to have the same dtype"): + copy_dst.index_copy_(0, index, torch.ones(2, dtype=torch.float64)) + + # Widening is rejected too, so ".to(dtype)" is required, not merely tidy. + widen_dst = torch.zeros(3, dtype=torch.float64) + with pytest.raises(RuntimeError, match="expected to have the same dtype"): + widen_dst.index_copy_(0, index, torch.ones(2, dtype=torch.float32)) + widen_dst.index_copy_(0, index, torch.ones(2, dtype=torch.float32).to(torch.float64)) + assert widen_dst.tolist() == [1.0, 1.0, 0.0] + + +# --------------------------------------------------------------------------- # +# T3 / T4 -- the two masked assignments that became torch.where. +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "fmax_values", + [ + [0.5, 2.0, 0.1, 3.0], + [9.0, 9.0, 9.0, 9.0], + [0.0, 0.0, 0.0, 0.0], + [1.0, float("nan"), 0.2, float("inf")], + ], +) +def test_smallest_fmax_where_matches_masked_assignment(fmax_values): + """The ``torch.where`` form of the ``smallest_fmax`` update is exact. + + Including NaN, which is the whole point. ``<`` is ``False`` for NaN, so the + original masked assignment *kept* the previous ``smallest_fmax`` when a + molecule's force went NaN. ``torch.minimum`` would instead propagate the NaN + and poison that molecule's oscillation tracking for the remainder of the + run. This test is what fails if someone "simplifies" the ``where`` into a + ``minimum``. + """ + fmax = torch.tensor(fmax_values) + smallest = torch.tensor([[1.0], [1.0], [1.0], [1.0]]) + + fmax_col = fmax.reshape(-1, 1) + reduced = (fmax_col < smallest).reshape(-1) + + reference = smallest.clone() + reference[reduced] = fmax_col[reduced] + rewrite = torch.where(reduced.unsqueeze(-1), fmax_col, smallest) + + assert torch.equal(rewrite, reference) + # And the guard against the tempting simplification: + if torch.isnan(fmax).any(): + assert not torch.equal(torch.minimum(fmax_col, smallest), reference) + + +@pytest.mark.parametrize( + "patterns", + [ + [[True, False, True, False]], + [[False, False, False, False]], + [[True, True, True, True]], + [[False, False, True, False], [False, True, False, False], [False] * 4], + ], +) +def test_oscillating_count_where_matches_masked_assignment(patterns): + """The fused ``torch.where`` counter update is integer-exact. + + The original was two statements -- zero the reduced entries, then increment + the non-reduced ones -- and the ordering mattered: a reduced molecule was + zeroed and then incremented by ``False``, landing on 0. The single ``where`` + (reduced -> 0, else -> old + 1) reproduces that, including across multiple + accumulating steps. + """ + reference = torch.zeros(4, dtype=torch.long) + rewrite = torch.zeros(4, dtype=torch.long) + + for pattern in patterns: + reduced = torch.tensor(pattern) + + reference[reduced] = 0 + reference += ~reduced + + rewrite = torch.where(reduced, torch.zeros_like(rewrite), rewrite + 1) + + assert torch.equal(rewrite, reference) + + +# --------------------------------------------------------------------------- # +# T5 -- the keystone: bit-identity against the old boolean-mask loop. +# --------------------------------------------------------------------------- # + + +def _reference_n_steps(state, n, opttol, patience, atom_mask=None): + """The pre-M6 boolean-mask loop, reimplemented here on purpose. + + This is a *reference*, not a copy of the code under test: it imports nothing + from ``optimization_engine`` and does all of its subsetting with boolean + masks, including inlining the old body of ``FIRE.clean`` (which now takes an + int64 index). Only ``FIRE.__call__`` is shared, because the FIRE step math + is not what M6 changed and is covered by ``test_fire_optimizer.py``. + + Keeping the reference in-process rather than checking in a golden + trajectory is deliberate: a golden file would compare float results across + platforms and BLAS versions, and would fail for reasons that have nothing + to do with this rewrite. + """ + numbers = state["numbers"] + coord = state["coord"] + if atom_mask is None: + atom_mask = torch.ones_like(numbers, dtype=torch.bool) + state["atom_mask"] = atom_mask + + optimizer = FIRE(coord) + smallest_fmax0 = torch.full((len(coord), 1), 999.0, dtype=torch.float, device=coord.device) + state["oscillating_count"] = torch.zeros(len(coord), dtype=torch.long, device=coord.device) + + istep = 0 + for istep in range(1, n + 1): + not_converged = ~state["converged_mask"] + if istep % 10 == 0 and not not_converged.any(): + break + + coord = state["coord"][not_converged] + if coord.shape[0] == 0: + break + numbers = state["numbers"][not_converged] + charges = state["charges"][not_converged] + atom_mask_subset = state["atom_mask"][not_converged] + smallest_fmax = smallest_fmax0[not_converged] + oscillating_count = state["oscillating_count"][not_converged] + + coord.requires_grad_(True) + energy, forces = state["nn"].forward_batched( + coord, numbers, charges, atom_mask=atom_mask_subset + ) + coord.requires_grad_(False) + + forces = forces.masked_fill(~atom_mask_subset.unsqueeze(-1), 0.0) + fmax = forces.norm(dim=-1).max(dim=-1)[0] + not_converged_post1 = fmax > opttol + stepped = optimizer(coord, forces).detach() + coord = torch.where(not_converged_post1.view(-1, 1, 1), stepped, coord) + + fmax_reduced = (fmax.reshape(-1, 1) < smallest_fmax).reshape(-1) + smallest_fmax[fmax_reduced] = fmax.reshape(-1, 1)[fmax_reduced] + oscillating_count[fmax_reduced] = 0 + oscillating_count += ~fmax_reduced + not_oscillating = oscillating_count < patience + not_converged_post = not_converged_post1 & not_oscillating + + # The old FIRE.clean body, verbatim, so the reference does not depend on + # the new int64 signature. + optimizer.v = optimizer.v[not_converged_post] + optimizer.Nsteps = optimizer.Nsteps[not_converged_post] + optimizer.dt = optimizer.dt[not_converged_post] + optimizer.a = optimizer.a[not_converged_post] + + state["converged_mask"][not_converged] = ~not_converged_post + state["fmax"][not_converged] = fmax + state["energy"][not_converged] = energy.detach().to(state["energy"].dtype) + state["coord"][not_converged] = coord + smallest_fmax0[not_converged] = smallest_fmax + state["oscillating_count"][not_converged] = oscillating_count + + final_coord = state["coord"].detach().clone().requires_grad_(True) + e_final, f_final = state["nn"].forward_batched( + final_coord, state["numbers"], state["charges"], atom_mask=state["atom_mask"] + ) + state["energy"] = e_final.detach().to(state["energy"].dtype) + f_final = f_final.detach().masked_fill(~state["atom_mask"].unsqueeze(-1), 0.0) + state["fmax"] = f_final.norm(dim=-1).max(dim=-1)[0].to(state["fmax"].dtype) + return istep + + +_STAGGERED_K = {i: 1.0 + 3.0 * i for i in range(8)} +_WIDE_K = {i: 0.5 + 0.37 * i for i in range(64)} + +# (label, nn factory, batch, natoms, n, opttol, patience, seed, atom_mask) +_SCENARIOS = [ + ("staggered convergence", lambda: _AnisotropicHarmonic(_STAGGERED_K), + 8, 6, 2000, 0.01, 5000, 0, None), + ("staggered plus oscillation drops", lambda: _AnisotropicHarmonic(_STAGGERED_K), + 8, 6, 2000, 1e-9, 20, 1, None), + ("all oscillating", _ConstantForce, 6, 5, 200, 0.01, 7, 2, None), + ("single molecule", lambda: _AnisotropicHarmonic(_STAGGERED_K), + 1, 4, 500, 0.01, 5000, 3, None), + ("n=0, loop never runs", lambda: _AnisotropicHarmonic(_STAGGERED_K), + 4, 4, 0, 0.01, 5000, 4, None), + ("large batch", lambda: _AnisotropicHarmonic(_WIDE_K), 64, 9, 800, 0.02, 300, 5, None), + ("padded batch, 2 ghost slots", lambda: _AnisotropicHarmonic(_STAGGERED_K), + 8, 6, 2000, 0.01, 5000, 6, "two-ghosts"), +] +_SCENARIOS += [ + (f"randomized seed {seed}", lambda: _AnisotropicHarmonic(_STAGGERED_K), + 8, 6, 400, 0.03, 60, seed, None) + for seed in range(7, 17) +] + +_STATE_KEYS = ("coord", "energy", "fmax", "converged_mask", "oscillating_count") + + +class TestEngineMatchesBooleanMaskReference: + """``n_steps`` is bit-identical to the boolean-mask loop it replaced.""" + + @pytest.mark.parametrize( + "label,nn_factory,batch,natoms,n,opttol,patience,seed,mask_kind", + _SCENARIOS, + ids=[s[0] for s in _SCENARIOS], + ) + def test_state_is_bit_identical(self, label, nn_factory, batch, natoms, n, + opttol, patience, seed, mask_kind): + """Every state tensor matches exactly -- ``torch.equal``, not ``allclose``. + + ``index_select(0, nonzero(m))`` and ``x[m]`` gather the same rows in the + same order, ``index_copy_`` and ``x[m] = v`` write the same rows, and no + arithmetic or reduction order changed. So the correct assertion is exact + equality; a tolerance would hide precisely the kind of reordering bug + this is meant to detect. + """ + atom_mask = None + if mask_kind == "two-ghosts": + atom_mask = torch.ones(batch, natoms, dtype=torch.bool) + atom_mask[:, -2:] = False + + reference_state = _make_state(nn_factory(), batch, natoms, seed) + production_state = _make_state(nn_factory(), batch, natoms, seed) + + _reference_n_steps( + reference_state, n=n, opttol=opttol, patience=patience, + atom_mask=None if atom_mask is None else atom_mask.clone(), + ) + n_steps( + production_state, n=n, opttol=opttol, patience=patience, + atom_mask=None if atom_mask is None else atom_mask.clone(), + ) + + mismatched = [ + key for key in _STATE_KEYS + if not torch.equal(reference_state[key], production_state[key]) + ] + assert not mismatched, ( + f"{label}: state differs from the boolean-mask reference in " + f"{mismatched}; max abs deltas " + + ", ".join( + f"{key}=" + f"{(reference_state[key].to(torch.double) - production_state[key].to(torch.double)).abs().max().item():.3e}" + for key in mismatched + ) + ) + + def test_scenarios_actually_exercise_partial_gathers(self): + """Guard the guard: the staggered scenario must converge unevenly. + + If every molecule converged on the same step, the subset gather would + always be the whole batch and the bit-identity above would be vacuous -- + it would never compare a *partial* gather, which is the only thing M6 + changed. Observe the width of the batch handed to the model each step: + a staggered run shrinks in several stages. + """ + class _WidthRecorder: + def __init__(self, inner): + self.inner = inner + self.widths: list[int] = [] + + def forward_batched(self, coord, numbers, charges, atom_mask=None): + self.widths.append(int(coord.shape[0])) + return self.inner.forward_batched(coord, numbers, charges, + atom_mask=atom_mask) + + recorder = _WidthRecorder(_AnisotropicHarmonic(_STAGGERED_K)) + state = _make_state(recorder, 8, 6, seed=0) + n_steps(state, n=2000, opttol=0.01, patience=5000) + + # The last call is the end-of-function recompute, which is always + # full-width; the loop body is everything before it. + in_loop = recorder.widths[:-1] + assert in_loop[0] == 8 + assert len(set(in_loop)) >= 3, ( + "the active set never shrank in stages, so no partial gather was " + f"ever exercised: observed batch widths {sorted(set(in_loop))}" + ) + assert min(in_loop) < 8 + + +class TestStepForStepIdentity: + """Trajectories agree at *every* prefix length, not just at the end. + + Final-state equality can in principle hide two errors that cancel. Comparing + after 1, 2, 3, ... steps localizes any divergence to the first step where it + appears, which is what makes the decomposition of ``n_steps`` into + ``_step_active_subset`` / ``_scatter_back`` / ``_recompute_final_energy_and_fmax`` + checkable rather than merely plausible: the refactor and the indexing change + landed together, so nothing may be taken on faith about either. + + ``n_steps`` builds a fresh ``FIRE`` sized to the full batch on entry, so it + cannot be driven one step at a time; each prefix length is a fresh run from + identical initial conditions instead. + """ + + @pytest.mark.parametrize("mask_kind", [None, "two-ghosts"]) + def test_every_prefix_of_the_trajectory_matches(self, mask_kind): + """40 prefixes x 5 state tensors, exact equality, on a hermetic potential.""" + batch, natoms = 8, 6 + atom_mask = None + if mask_kind == "two-ghosts": + atom_mask = torch.ones(batch, natoms, dtype=torch.bool) + atom_mask[:, -2:] = False + + for steps in range(1, 41): + reference_state = _make_state( + _AnisotropicHarmonic(_STAGGERED_K), batch, natoms, seed=0) + production_state = _make_state( + _AnisotropicHarmonic(_STAGGERED_K), batch, natoms, seed=0) + _reference_n_steps( + reference_state, n=steps, opttol=0.05, patience=12, + atom_mask=None if atom_mask is None else atom_mask.clone()) + n_steps( + production_state, n=steps, opttol=0.05, patience=12, + atom_mask=None if atom_mask is None else atom_mask.clone()) + + for key in _STATE_KEYS: + assert torch.equal(reference_state[key], production_state[key]), ( + f"diverged at step {steps} in {key!r}" + ) + + def test_the_prefixes_are_not_all_the_same_state(self): + """Guard the guard: the trajectory must actually move across prefixes. + + If every prefix produced identical state (already converged at step 1), + the 40 comparisons above would be 40 copies of one comparison. + """ + seen = set() + for steps in (1, 5, 20, 40): + state = _make_state(_AnisotropicHarmonic(_STAGGERED_K), 8, 6, seed=0) + n_steps(state, n=steps, opttol=0.05, patience=12) + seen.add(tuple(state["coord"].reshape(-1).tolist())) + assert len(seen) == 4 + + +class TestDecomposedHelpers: + """The four pieces ``n_steps`` was split into are callable on their own. + + The point of the split is not line count -- it is that each piece can be + driven directly with a state it did not build itself, which the 200-line + original could not be. + """ + + def _state(self): + state = _make_state(_AnisotropicHarmonic(_STAGGERED_K), 4, 5, seed=0) + state["atom_mask"] = torch.ones_like(state["numbers"], dtype=torch.bool) + state["oscillating_count"] = torch.zeros(4, dtype=torch.long) + return state + + def test_step_active_subset_returns_rows_aligned_with_the_index(self): + """Every field of the result has the active subset's leading dimension.""" + from Auto3D.batch_opt.optimization_engine import _step_active_subset + + state = self._state() + optimizer = FIRE(state["coord"].index_select(0, torch.tensor([1, 3]))) + smallest = torch.full((4, 1), 999.0) + active = torch.tensor([1, 3]) + + result = _step_active_subset(state, optimizer, active, smallest, + opttol=0.01, patience=100) + + assert result.coord.shape == (2, 5, 3) + for field in (result.energy, result.fmax, result.still_active, + result.oscillating_count): + assert field.shape[0] == 2 + assert result.smallest_fmax.shape == (2, 1) + # And it left the full-batch state alone: writing is _scatter_back's job. + assert torch.equal(state["converged_mask"], torch.zeros(4, dtype=torch.bool)) + + def test_scatter_back_writes_only_the_active_rows(self): + """Inactive rows keep their previous values, exactly. + + A scatter that touched an untouched row would silently overwrite a + converged structure's final geometry with a stale one. + """ + from Auto3D.batch_opt.optimization_engine import _StepResult, _scatter_back + + state = self._state() + state["coord"] = torch.zeros(4, 5, 3) + state["fmax"] = torch.full((4,), 7.0) + state["energy"] = torch.full((4,), 3.0, dtype=torch.double) + smallest = torch.full((4, 1), 999.0) + active = torch.tensor([1, 3]) + + _scatter_back(state, active, smallest, _StepResult( + coord=torch.ones(2, 5, 3), + energy=torch.full((2,), -1.0, dtype=torch.double), + fmax=torch.full((2,), 0.5), + still_active=torch.tensor([True, False]), + smallest_fmax=torch.full((2, 1), 0.5), + oscillating_count=torch.tensor([0, 4]), + )) + + assert state["converged_mask"].tolist() == [False, False, False, True] + assert state["fmax"].tolist() == [7.0, 0.5, 7.0, 0.5] + assert state["energy"].tolist() == [3.0, -1.0, 3.0, -1.0] + assert torch.equal(state["coord"][0], torch.zeros(5, 3)) + assert torch.equal(state["coord"][1], torch.ones(5, 3)) + assert smallest.reshape(-1).tolist() == [999.0, 0.5, 999.0, 0.5] + assert state["oscillating_count"].tolist() == [0, 0, 0, 4] + + def test_scatter_back_casts_to_the_destination_dtype(self): + """float64 sources land in float32 destinations without raising. + + ``index_copy_`` refuses to cast in either direction, so the cast has to be + at the call site. This is the unit-level statement of what + ``test_float64_model_outputs_do_not_raise`` checks end to end. + """ + from Auto3D.batch_opt.optimization_engine import _StepResult, _scatter_back + + state = self._state() + smallest = torch.full((4, 1), 999.0) + _scatter_back(state, torch.tensor([0]), smallest, _StepResult( + coord=torch.ones(1, 5, 3, dtype=torch.float64), + energy=torch.zeros(1, dtype=torch.float32), + fmax=torch.full((1,), 0.25, dtype=torch.float64), + still_active=torch.tensor([True]), + smallest_fmax=torch.full((1, 1), 0.25, dtype=torch.float64), + oscillating_count=torch.tensor([2], dtype=torch.int32), + )) + assert state["coord"].dtype is torch.float32 + assert state["fmax"].dtype is torch.float32 + assert state["energy"].dtype is torch.float64 + assert state["oscillating_count"].dtype is torch.long + assert state["fmax"][0].item() == 0.25 + + def test_emit_progress_is_a_no_op_without_a_callback(self): + """No callback means no ``optimization_counts``, hence no sync at all.""" + from Auto3D.batch_opt.optimization_engine import _emit_progress + + state = self._state() + counter = SyncCounter() + with counter: + _emit_progress(state, patience=100, progress_cb=None, istep=5) + assert counter.total == 0, counter.report() + + def test_emit_progress_swallows_a_failing_callback(self): + """A broken progress display must never abort an optimization.""" + from Auto3D.batch_opt.optimization_engine import _emit_progress + + def explode(event): + raise RuntimeError("display is on fire") + + _emit_progress(self._state(), patience=100, progress_cb=explode, istep=5) + + def test_emit_progress_reports_the_counts(self): + """The event carries the five documented keys.""" + from Auto3D.batch_opt.optimization_engine import _emit_progress + + state = self._state() + state["converged_mask"] = torch.tensor([True, True, False, False]) + events: list[dict] = [] + _emit_progress(state, patience=100, progress_cb=events.append, istep=7) + assert events == [{"step": 7, "total": 4, "converged": 2, "dropped": 0, + "active": 2}] + + def test_recompute_final_energy_and_fmax_uses_the_stored_coordinates(self): + """Reported energy/fmax describe the reported geometry, not the pre-step one.""" + from Auto3D.batch_opt.optimization_engine import ( + _recompute_final_energy_and_fmax, + ) + + state = self._state() + state["coord"] = torch.full((4, 5, 3), 0.5) + state["energy"] = torch.full((4,), 12345.0, dtype=torch.double) + state["fmax"] = torch.full((4,), 999.0) + + _recompute_final_energy_and_fmax(state) + + expected_energy = torch.tensor( + [_STAGGERED_K[i] * 5 * 3 * 0.25 for i in range(4)], dtype=torch.double) + assert torch.allclose(state["energy"], expected_energy) + expected_fmax = torch.tensor( + [(2.0 * _STAGGERED_K[i] * 0.5) * (3 ** 0.5) for i in range(4)]) + assert torch.allclose(state["fmax"], expected_fmax, rtol=1e-5) + + def test_recompute_ignores_padded_atom_forces(self): + """Ghost slots cannot inflate the reported fmax.""" + from Auto3D.batch_opt.optimization_engine import ( + _recompute_final_energy_and_fmax, + ) + + state = self._state() + state["coord"] = torch.full((4, 5, 3), 0.1) + state["coord"][:, -1, :] = 50.0 # a wildly displaced ghost atom + state["atom_mask"][:, -1] = False + + _recompute_final_energy_and_fmax(state) + + quiet = torch.tensor( + [(2.0 * _STAGGERED_K[i] * 0.1) * (3 ** 0.5) for i in range(4)]) + assert torch.allclose(state["fmax"], quiet, rtol=1e-5) + + +# --------------------------------------------------------------------------- # +# T6 -- the regression lock. +# --------------------------------------------------------------------------- # + + +def _loop_body_syncs(steps: int) -> SyncCounter: + """Count syncs for exactly ``steps`` full-width loop iterations. + + ``patience`` is huge and the force is constant, so nothing ever converges + and every step runs the full body. ``steps < 10`` keeps the throttled + ``not_converged.any()`` (every 10 steps) and ``print_stats`` (``n >= 10``) + off the hot path, so the count is the loop body alone plus a fixed + end-of-function constant -- which cancels in the delta below. + """ + counter = SyncCounter(attribute=True) + state = _make_state(_ConstantForce(), 8, 6, seed=0) + with counter: + n_steps(state, n=steps, opttol=0.0, patience=10 ** 9) + return counter + + +class TestHotLoopDoesNotSync: + """The counted, CI-enforceable half of this cluster.""" + + def test_hot_loop_does_no_boolean_mask_indexing(self): + """Zero boolean-mask reads and zero boolean-mask writes in ``n_steps``. + + This fails the moment someone writes ``state['coord'][mask]`` again. + It is the durable part of the change: CI cannot time the loop, but it + can prove the sync-forcing ops are gone. + """ + counter = _loop_body_syncs(9) + assert counter.bool_mask_ops == 0, ( + "n_steps performed boolean-mask indexing, which forces a GPU->CPU " + "sync per call. Use index_select / index_copy_ with the int64 index " + f"from the single torch.nonzero instead.\n{counter.report()}" + ) + + def test_hot_loop_syncs_at_most_twice_per_step(self): + """At most 2 sync-forcing ops per step, and they are ``nonzero`` calls. + + Two, not zero: one ``nonzero`` for the active subset (reused by six + gathers and six scatters) and one for ``FIRE.clean``, whose mask is + indexed within the active subset rather than the full batch. The delta + between two step counts cancels every fixed end-of-function cost. + """ + few, many = 4, 9 + delta = _loop_body_syncs(many).total - _loop_body_syncs(few).total + per_step = delta / (many - few) + assert per_step <= 2.0, ( + f"{per_step:.1f} sync-forcing ops per optimization step, expected " + f"<= 2.0\n{_loop_body_syncs(many).report()}" + ) + + def test_the_two_remaining_syncs_are_both_nonzero(self): + """Name the survivors, so the accounting stays honest. + + ``nonzero`` *is* a sync on CUDA -- it is the mechanism by which + boolean-mask indexing syncs in the first place. The win is that one + ``nonzero`` result is reused twelve times instead of each gather and + scatter computing its own, i.e. 18 -> 2, never 18 -> 0. + """ + counter = _loop_body_syncs(9) + non_nonzero = { + label: count for label, count in counter.counts.items() + if label != NONZERO and count + } + # print_stats runs once at the end of n_steps and reads two scalars. + readbacks = sum(non_nonzero.values()) + assert readbacks <= 2, ( + "unexpected sync-forcing ops beyond the two nonzero calls and the " + f"final print_stats: {non_nonzero}\n{counter.report()}" + ) + assert counter.counts[NONZERO] == 2 * 9, ( + f"expected 2 nonzero calls per step over 9 steps, got " + f"{counter.counts[NONZERO]}\n{counter.report()}" + ) + + def test_counter_would_catch_a_reintroduced_boolean_mask(self): + """Prove the detector detects. Otherwise T6 could pass vacuously. + + A counter that silently stopped classifying ``aten.index.Tensor`` would + make every assertion above trivially true, so exercise it on a known + boolean-mask read and a known masked write. + """ + counter = SyncCounter() + with counter: + x = torch.randn(4, 3) + mask = torch.tensor([True, False, True, False]) + _ = x[mask] + x[mask] = torch.zeros(2, 3) + assert counter.bool_mask_ops == 2, counter.report() + assert set(counter.counts) <= set(BOOL_MASK_LABELS) | {NONZERO} + + +# --------------------------------------------------------------------------- # +# The dtype hazard index_copy_ introduces, end to end. +# --------------------------------------------------------------------------- # + + +class _Float64Model: + """A model returning float64 energies *and* float64 forces. + + This is the realistic custom-NNP case, and on the pre-M6 loop it **crashed**: + ``smallest_fmax`` is allocated float32, ``fmax`` inherits float64 from the + forces, and ``smallest_fmax[fmax_reduced] = fmax.reshape(-1, 1)[fmax_reduced]`` + raised ``"Index put requires the source and destination dtypes match"`` + whenever two or more molecules reduced their force in the same step. With + exactly one reducing molecule the value had ``numel() == 1``, took ATen's + ``masked_fill_`` fast path, and silently cast -- so the failure was + *batch-size dependent* and invisible in a single-molecule test. + + Casting explicitly at every ``index_copy_`` site fixes that as a side effect. + This is the one place the rewrite is deliberately **not** bit-identical to + its predecessor: it succeeds where the predecessor raised. + """ + + def forward_batched(self, coord, numbers, charges, atom_mask=None): + energy = (coord.to(torch.float64) ** 2).sum(dim=(1, 2)) + forces = -2.0 * coord.to(torch.float64) + return energy, forces + + +def test_float64_model_outputs_do_not_raise(): + """A float64 NNP runs to completion and ``state`` dtypes are preserved.""" + state = _make_state(_Float64Model(), 4, 5, seed=0) + n_steps(state, n=400, opttol=0.01, patience=5000) + + assert state["coord"].dtype is torch.float32 + assert state["fmax"].dtype is torch.float32 + assert state["energy"].dtype is torch.float64 + assert state["oscillating_count"].dtype is torch.long + assert torch.isfinite(state["coord"]).all() + # Not merely "did not raise": the run must actually reach the minimum, so + # the multi-molecule force-reduction path (the one that used to raise) is + # genuinely exercised many times over. + assert state["converged_mask"].all() diff --git a/tests/test_parallel_embed.py b/tests/test_parallel_embed.py index 40b5e1a6..c4e273b1 100644 --- a/tests/test_parallel_embed.py +++ b/tests/test_parallel_embed.py @@ -1,5 +1,5 @@ # tests/test_parallel_embed.py -"""Tests for parallel conformer embedding module.""" +"""Tests for Auto3D.embedding (parallel conformer embedding).""" import multiprocessing as mp import os from concurrent.futures.process import BrokenProcessPool @@ -7,7 +7,7 @@ import pytest from rdkit import Chem -from Auto3D.isomers.parallel_embed import _embed_single, embed_conformers_parallel +from Auto3D.embedding import _embed_single, embed_conformers_parallel def _suicide_embed(smi, name, n_conformers, threshold, np_threads): @@ -49,7 +49,7 @@ def test_embed_single_with_dynamic_conformers(self): molecule, where embedding + RMSD pruning collapses to 1 regardless of the requested count and could hide a formula regression). """ - from Auto3D.utils.chemistry import calculate_conformer_count + from Auto3D.utils.molprops import calculate_conformer_count mol = Chem.AddHs(Chem.MolFromSmiles("CCCCCC")) # hexane: flexible expected_upper_bound = calculate_conformer_count(mol) @@ -242,7 +242,7 @@ def test_parallel_embed_reraises_broken_pool(self, monkeypatch): # Replace the worker with one that kills its process mid-task. monkeypatch.setattr( - "Auto3D.isomers.parallel_embed._embed_single", _suicide_embed + "Auto3D.embedding._embed_single", _suicide_embed ) with pytest.raises(BrokenProcessPool): diff --git a/tests/test_pipeline_e2e.py b/tests/test_pipeline_e2e.py index a8b7cdf4..5621399e 100644 --- a/tests/test_pipeline_e2e.py +++ b/tests/test_pipeline_e2e.py @@ -263,7 +263,7 @@ class TestClashReliefWarning: vanishes from the output with no other trace. Before this test, ``grep -rn "produced no conformers after clash relief" tests/`` returned nothing: this is distinct from (and untested by) the parallel-embedding - path's own version of the same warning in ``isomers/parallel_embed.py``, + path's own version of the same warning in ``embedding.py``, which ``test_workflow.py`` already covers. Hermetic: no NNP, no network. ``relieve_clash`` itself is monkeypatched, diff --git a/tests/test_processors.py b/tests/test_processors.py index f0c98881..25440b9c 100644 --- a/tests/test_processors.py +++ b/tests/test_processors.py @@ -79,7 +79,14 @@ def test_rdkit_engine(self, tmp_path): processor = TautomerProcessor(config) result = processor.process(str(input_file), str(output_file)) - assert Path(result).exists() + # `.exists()` alone is satisfied by an empty file, which is also what + # the *disabled* branch's contract looks like from the outside if + # `result` were accidentally the input path. Pin both: the enabled + # branch must return the OUTPUT path specifically, and that file must + # actually carry the tautomer engine's records, not be empty. + assert result == str(output_file) + written = Path(result).read_text().strip() + assert written, "tautomer engine produced no records in the output file" def test_tautomer_processor_uses_facade(monkeypatch, tmp_path): @@ -95,11 +102,22 @@ def fake_create(engine_type, input_path, output_path, pka_norm=True): monkeypatch.setattr(proc, "create_tautomer_engine", fake_create) monkeypatch.setattr(proc, "hash_taut_smi", lambda a, b: None) - cfg = Auto3DOptions(path="x.smi", k=1, enumerate_tautomer=True, tauto_engine="rdkit") + # pKaNorm=False is deliberately non-default (Auto3DOptions defaults it to + # True): a mutation that hardcodes the forwarded value instead of reading + # self.config.pKaNorm would otherwise still match a default-valued cfg. + cfg = Auto3DOptions( + path="x.smi", k=1, enumerate_tautomer=True, tauto_engine="rdkit", + pKaNorm=False, + ) out = proc.TautomerProcessor(cfg).process("in.smi", "out.smi") assert out == "out.smi" assert calls["ran"] is True - assert calls["args"][0] == "rdkit" + # Pin the full argument tuple `create_tautomer_engine` was called with, not + # just the engine name: `args[1]`/`args[2]` catch input_path/output_path + # being swapped or dropped, and `args[3]` catches the wrong (or a missing) + # pKaNorm being forwarded -- none of which the old, name-only assertion + # could have caught. + assert calls["args"] == ("rdkit", "in.smi", "out.smi", False) def test_tautomer_processor_skips_when_disabled(): diff --git a/tests/test_ranking.py b/tests/test_ranking.py index b690f08d..de7be5f3 100644 --- a/tests/test_ranking.py +++ b/tests/test_ranking.py @@ -294,7 +294,7 @@ def test_top_k_equals_1_skips_broken_connectivity(self, tmp_path): mol_valid = _create_mol_with_energy("CC", -9.0, "mol") # Sanity: the broken one fails check_connectivity, the valid one passes. - from Auto3D.utils.chemistry import check_connectivity + from Auto3D.utils.connectivity import check_connectivity assert check_connectivity(mol_broken) is False assert check_connectivity(mol_valid) is True @@ -327,7 +327,7 @@ def test_top_k_equals_1_returns_empty_when_all_broken(self, tmp_path): pos = conf.GetAtomPosition(0) conf.SetAtomPosition(0, (pos.x + 5.0, pos.y, pos.z)) - from Auto3D.utils.chemistry import check_connectivity + from Auto3D.utils.connectivity import check_connectivity assert check_connectivity(mol_broken) is False ranker = ConformerRanker( @@ -719,7 +719,7 @@ def _canonical(mol: Chem.Mol) -> str: def test_enumerate_isomer_false_returns_both_molecules(self, tmp_path): from Auto3D.isomer_engine import RDKitIsomer from Auto3D.ranking import ConformerRanker, species_id - from Auto3D.utils.file_ops import smiles2smi + from Auto3D.utils.smi_io import smiles2smi smi_path = str(tmp_path / "in.smi") smiles2smi([self.PYRIDONE, self.HYDROXYPYRIDINE], smi_path) diff --git a/tests/test_stereo_identity.py b/tests/test_stereo_identity.py index f29736fd..383ec1e3 100644 --- a/tests/test_stereo_identity.py +++ b/tests/test_stereo_identity.py @@ -101,8 +101,8 @@ def test_unspecified_center_is_enumerated_or_refused(self, job_dir): This writes a genuine flat (2D, no wedge bonds, no parity flags) SDF record for alanine to disk and feeds it through the production - ``rdkit_sdf`` engine -- the same ``RDKitSdfIsomerAdapter`` / - ``RDKitSdfIsomer.run()`` the pipeline dispatches to for SDF input -- + ``rdkit_sdf`` engine -- the same ``RDKitSdfIsomer.run()`` the pipeline + dispatches to for SDF input -- via ``Auto3D.isomers.factory.create_isomer_engine``. It then inspects the SDF file Auto3D actually writes, grouped by species name (the conformer-index suffix stripped). Either the two configurations must diff --git a/tests/test_stereo_postopt.py b/tests/test_stereo_postopt.py index d11e661e..1bee6be3 100644 --- a/tests/test_stereo_postopt.py +++ b/tests/test_stereo_postopt.py @@ -14,7 +14,7 @@ from Auto3D.filtering import filter_unique_optimized from Auto3D.ranking import ConformerRanker -from Auto3D.utils.chemistry import filter_unique +from Auto3D.filtering import filter_unique from Auto3D.utils.stereo_check import ( STEREO_CHANGED_PROP, apply_optimized_coords, diff --git a/tests/test_thermo_imaginary_mode_inversion.py b/tests/test_thermo_imaginary_mode_inversion.py index 7f25d618..2d40c7f0 100644 --- a/tests/test_thermo_imaginary_mode_inversion.py +++ b/tests/test_thermo_imaginary_mode_inversion.py @@ -36,13 +36,14 @@ from ase.vibrations import VibrationsData import Auto3D.ASE.thermo as thermo_mod -from Auto3D.ASE.thermo import analyze_vibrations, ev2hatree, projected_vibrations +from Auto3D.ASE.thermo import analyze_vibrations, projected_vibrations from Auto3D.constants import ( EV_PER_WAVENUMBER, + EV_TO_HARTREE, IMAGINARY_MODE_CUTOFF_CM, LOW_FREQUENCY_CUTOFF_CM, ) -from Auto3D.utils.chemistry import EV_TO_KCAL_PER_MOL, HARTREE_TO_KCAL_PER_MOL +from Auto3D.utils.energy import EV_TO_KCAL_PER_MOL, HARTREE_TO_KCAL_PER_MOL from tests.helpers_vibrations import ( ASE_SELECTION_RULES, atoms_for, @@ -116,7 +117,7 @@ def _gibbs_kcal(atoms, modes, potential_energy=0.0) -> float: thermo.get_gibbs_energy( temperature=T_REFERENCE, pressure=PRESSURE_PA, verbose=False ) - * ev2hatree + * EV_TO_HARTREE * HARTREE_TO_KCAL_PER_MOL ) diff --git a/tests/test_thermo_reference.py b/tests/test_thermo_reference.py index 482d99c2..89e4d9d2 100644 --- a/tests/test_thermo_reference.py +++ b/tests/test_thermo_reference.py @@ -61,8 +61,20 @@ def test_malformed_record_does_not_abort_the_batch(self, job_dir): out = calc_thermo(str(combined), "AIMNET", use_gpu=False) results = [m for m in Chem.SDMolSupplier(str(out), removeHs=False) if m] - assert any(m.HasProp("G_hartree") for m in results), ( - "the valid molecules produced no thermo result" + # `any(...)` passes even if only one of the two valid records survived + # -- e.g. a regression that aborts the batch partway through, right + # after the corrupt record, would still satisfy it. Both good records + # must independently reach the output with a Gibbs energy, and the + # corrupt one must not have been counted or duplicated into a third + # record. + with_g = { + m.GetProp("_Name") for m in results if m.HasProp("G_hartree") + } + assert with_g == {"ethanol", "propanol"}, ( + f"expected both valid molecules to produce a thermo result, got {with_g}" + ) + assert len(results) == 2, ( + f"expected exactly the two valid records, got {len(results)}" ) @@ -70,22 +82,27 @@ class TestStationaryPointGating: """G must not be reported for a structure the optimizer did not converge.""" def test_unconverged_geometry_is_flagged_or_refused(self, job_dir): - """With opt_steps=1 nothing can converge, so no G may be emitted unflagged. - - The non-vacuity check at the end of this test previously asserted - `with_g` directly, which silently assumed the flag-and-emit - resolution (G is still reported, just marked approximate). Task 5 - implemented the other resolution this test's own name and the - spec's exit criterion both allow -- refusing to emit G at all for a - structure that did not converge (`mol.SetProp("Thermo_failed", - "not_converged")`, no `G_hartree` set) -- so `with_g` is legitimately - empty for a single-molecule input under that resolution, and the old - assertion failed for a reason unrelated to M8. Non-vacuity is - re-established below without assuming which resolution was taken: by - confirming the run produced output at all, and that the input - molecule is accounted for either way -- emitted with a flagged G, or - emitted carrying a failure marker (`Thermo_failed`, the resolution - this codebase actually implements) instead of one. + """With opt_steps=1 nothing can converge, so no G may be emitted at all. + + This test used to check `mol.HasProp("Thermo_converged") or + mol.HasProp("Thermo_warning")` for the (hypothetical) case where a + structure that failed to converge still got a flagged G -- but + neither property is ever set anywhere in the source tree, so that + loop could only fail, never pass on its own merits; it only appeared + to pass because `with_g` is always empty under the resolution this + codebase actually implements (`calc_thermo`'s stationary-point gate + withholds G entirely and sets `Thermo_failed="not_converged"` + instead). The final non-vacuity check was equally dead: every record + `_write_thermo_output` writes -- success or failure -- always ends up + with `Thermo_failed` set (successes get `""` if not already set, + failures get it set explicitly before being appended), so + `m.HasProp("G_hartree") or m.HasProp("Thermo_failed")` was true for + any record calc_thermo could possibly write, regardless of what it + actually computed. + + Replaced with an assertion of the one behavior calc_thermo + implements: no G_hartree anywhere, and the record explicitly marked + `Thermo_failed == "not_converged"`. """ from Auto3D.ASE.thermo import calc_thermo @@ -101,21 +118,17 @@ def test_unconverged_geometry_is_flagged_or_refused(self, job_dir): out = calc_thermo(str(path), "AIMNET", opt_steps=1, use_gpu=False) results = [m for m in Chem.SDMolSupplier(str(out), removeHs=False) if m] - with_g = [m for m in results if m.HasProp("G_hartree")] - for mol in with_g: - assert mol.HasProp("Thermo_converged") or mol.HasProp("Thermo_warning"), ( - "G was reported for a structure that could not have converged in " - "one step, with no flag distinguishing it" - ) - - # Non-vacuity without assuming which resolution was taken: the run - # must have produced output at all, and the single input molecule - # must be accounted for either way -- emitted with a flagged G, or - # emitted carrying a failure marker instead of one. assert results, "calc_thermo produced no output records at all" + + with_g = [m for m in results if m.HasProp("G_hartree")] + assert not with_g, ( + "G_hartree was reported for a structure that could not have " + "converged in a single BFGS step; calc_thermo's only implemented " + "resolution for an unconverged geometry is to withhold G entirely" + ) assert all( - m.HasProp("G_hartree") or m.HasProp("Thermo_failed") for m in results - ), "a record carried neither a Gibbs energy nor a failure marker" + m.GetProp("Thermo_failed") == "not_converged" for m in results + ), "the unconverged record was not marked Thermo_failed='not_converged'" class TestHessianGeometry: diff --git a/tests/test_utils.py b/tests/test_utils.py index 201a30d2..7282e544 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,9 +3,9 @@ from rdkit import Chem import Auto3D from Auto3D.config import Auto3DOptions -from Auto3D.utils.chemistry import check_connectivity +from Auto3D.utils.connectivity import check_connectivity from Auto3D.utils.validation import check_input -from Auto3D.utils.file_ops import find_smiles_not_in_sdf +from Auto3D.utils.reconciliation import find_smiles_not_in_sdf folder = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) diff --git a/tests/test_utils_chemistry.py b/tests/test_utils_chemistry.py deleted file mode 100644 index 4c86adf4..00000000 --- a/tests/test_utils_chemistry.py +++ /dev/null @@ -1,723 +0,0 @@ -#!/usr/bin/env python -"""Tests for Auto3D.utils.chemistry module.""" -from __future__ import annotations - -import numpy as np -import pytest -from rdkit import Chem -from rdkit.Chem import AllChem - -from Auto3D.utils.chemistry import ( - HARTREE_TO_EV, - HARTREE_TO_KCAL_PER_MOL, - EV_TO_KCAL_PER_MOL, - hartree2ev, - hartree2kcalpermol, - ev2kcalpermol, - get_mol_charge, - min_pairwise_distance, - get_rmsd, - check_connectivity, -) - - -class TestEnergyConversionConstants: - """Test energy conversion constants and aliases.""" - - def test_hartree_to_ev_value(self): - """Test that HARTREE_TO_EV has the correct CODATA 2018 value.""" - assert abs(HARTREE_TO_EV - 27.211386245988) < 1e-10 - - def test_hartree_to_kcal_per_mol_value(self): - """Test that HARTREE_TO_KCAL_PER_MOL has the expected value.""" - assert abs(HARTREE_TO_KCAL_PER_MOL - 627.50947337481) < 1e-8 - - def test_ev_to_kcal_per_mol_value(self): - """Test that EV_TO_KCAL_PER_MOL has the expected value.""" - assert abs(EV_TO_KCAL_PER_MOL - 23.060547830619026) < 1e-10 - - def test_backward_compatibility_aliases(self): - """Test that backward compatibility aliases match constants.""" - assert hartree2ev == HARTREE_TO_EV - assert hartree2kcalpermol == HARTREE_TO_KCAL_PER_MOL - assert ev2kcalpermol == EV_TO_KCAL_PER_MOL - - def test_conversion_consistency(self): - """Test that conversion factors are mathematically consistent.""" - # HARTREE_TO_KCAL_PER_MOL should approximately equal - # HARTREE_TO_EV * EV_TO_KCAL_PER_MOL - calculated = HARTREE_TO_EV * EV_TO_KCAL_PER_MOL - # Allow some tolerance for floating point precision - assert abs(calculated - HARTREE_TO_KCAL_PER_MOL) < 0.001 - - -class TestGetMolCharge: - """Test the get_mol_charge function.""" - - def test_neutral_molecule(self): - """Test charge of a neutral molecule.""" - mol = Chem.MolFromSmiles("CCO") - assert get_mol_charge(mol) == 0 - - def test_cation(self): - """Test charge of a cation.""" - mol = Chem.MolFromSmiles("[NH4+]") - assert get_mol_charge(mol) == 1 - - def test_anion(self): - """Test charge of an anion.""" - mol = Chem.MolFromSmiles("[O-]") - assert get_mol_charge(mol) == -1 - - def test_doubly_charged_cation(self): - """Test charge of a doubly charged cation.""" - mol = Chem.MolFromSmiles("[Ca+2]") - assert get_mol_charge(mol) == 2 - - def test_zwitterion(self): - """Test charge of a zwitterion (net neutral).""" - # Glycine zwitterion - mol = Chem.MolFromSmiles("[NH3+]CC([O-])=O") - assert get_mol_charge(mol) == 0 - - def test_multiple_charges(self): - """Test molecule with multiple charged atoms.""" - mol = Chem.MolFromSmiles("[O-]C([O-])=O") # Carbonate - assert get_mol_charge(mol) == -2 - - -class TestMinPairwiseDistance: - """Test the min_pairwise_distance function.""" - - def test_simple_three_points(self): - """Test with three simple points.""" - points = np.array([ - [0.0, 0.0, 0.0], - [1.0, 0.0, 0.0], - [0.0, 2.0, 0.0] - ]) - result = min_pairwise_distance(points) - assert abs(result - 1.0) < 1e-5 - - def test_two_points(self): - """Test with two points.""" - points = np.array([ - [0.0, 0.0, 0.0], - [3.0, 4.0, 0.0] # Distance = 5 - ]) - result = min_pairwise_distance(points) - assert abs(result - 5.0) < 1e-5 - - def test_collinear_points(self): - """Test with collinear points.""" - points = np.array([ - [0.0, 0.0, 0.0], - [1.0, 0.0, 0.0], - [2.0, 0.0, 0.0], - [5.0, 0.0, 0.0] - ]) - result = min_pairwise_distance(points) - assert abs(result - 1.0) < 1e-5 - - def test_3d_points(self): - """Test with points in 3D space.""" - points = np.array([ - [0.0, 0.0, 0.0], - [1.0, 1.0, 1.0], # Distance = sqrt(3) ~ 1.732 - [5.0, 5.0, 5.0] - ]) - result = min_pairwise_distance(points) - expected = np.sqrt(3) - assert abs(result - expected) < 1e-5 - - def test_input_type_conversion(self): - """Test that integer input is properly converted.""" - points = np.array([ - [0, 0, 0], - [1, 0, 0], - [0, 2, 0] - ], dtype=np.int32) - result = min_pairwise_distance(points) - assert abs(result - 1.0) < 1e-5 - - def test_very_close_points(self): - """Test with very close points.""" - points = np.array([ - [0.0, 0.0, 0.0], - [0.001, 0.0, 0.0], - [10.0, 0.0, 0.0] - ]) - result = min_pairwise_distance(points) - assert abs(result - 0.001) < 1e-6 - - -class TestGetRmsd: - """Test the get_rmsd function.""" - - def test_identical_molecules(self): - """Test RMSD of a molecule with itself is 0.""" - mol = Chem.MolFromSmiles("CCO") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - - mol_copy = Chem.Mol(mol) - rmsd = get_rmsd(mol, mol_copy) - assert abs(rmsd) < 1e-5 - - def test_different_conformers(self): - """Test RMSD of different conformers is > 0.""" - mol = Chem.MolFromSmiles("CCCCCC") # Hexane - flexible - mol = Chem.AddHs(mol) - - # Generate two different conformers - AllChem.EmbedMolecule(mol, randomSeed=42) - conf1 = mol.GetConformer() - - mol2 = Chem.Mol(mol) - AllChem.EmbedMolecule(mol2, randomSeed=123) - - rmsd = get_rmsd(mol, mol2) - # Different random seeds should give different conformers - # The RMSD should be >= 0 (might still be 0 if conformers happen to be similar) - assert rmsd >= 0 - - def test_remove_hs_option(self): - """Test that remove_hs option works.""" - mol = Chem.MolFromSmiles("C") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - - mol_copy = Chem.Mol(mol) - - # Both should return 0 for identical molecules - rmsd_with_hs_removed = get_rmsd(mol, mol_copy, remove_hs=True) - rmsd_without_hs_removed = get_rmsd(mol, mol_copy, remove_hs=False) - - assert abs(rmsd_with_hs_removed) < 1e-5 - assert abs(rmsd_without_hs_removed) < 1e-5 - - def test_mismatched_molecules_returns_inf(self): - """Mismatched molecules are incomparable and return float('inf'). - - An incomparable pair is treated as "distinct" (inf), matching - filter_unique, so a downstream `rmsd < threshold` check keeps the - structure instead of dropping it as a false duplicate. - """ - mol1 = Chem.MolFromSmiles("CCO") - mol1 = Chem.AddHs(mol1) - AllChem.EmbedMolecule(mol1, randomSeed=42) - - mol2 = Chem.MolFromSmiles("CCCC") - mol2 = Chem.AddHs(mol2) - AllChem.EmbedMolecule(mol2, randomSeed=42) - - # This should raise RuntimeError internally and return inf. - rmsd = get_rmsd(mol1, mol2) - assert rmsd == float("inf") - - def test_runtime_error_returns_inf(self, monkeypatch): - """A RuntimeError from the RMSD computation returns float('inf').""" - from Auto3D.utils import chemistry as chem - - mol = Chem.MolFromSmiles("CCO") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - mol_copy = Chem.Mol(mol) - - def _boom(*args, **kwargs): - raise RuntimeError("forced failure") - - monkeypatch.setattr(chem.rdMolAlign, "GetBestRMS", _boom) - assert chem.get_rmsd(mol, mol_copy) == float("inf") - - -class TestCheckConnectivity: - """Test the check_connectivity function.""" - - def test_valid_ethanol_connectivity(self): - """Test that a valid ethanol conformer has correct connectivity.""" - mol = Chem.MolFromSmiles("CCO") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - - assert check_connectivity(mol) is True - - def test_valid_methane_connectivity(self): - """Test that a valid methane conformer has correct connectivity.""" - mol = Chem.MolFromSmiles("C") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - - assert check_connectivity(mol) is True - - def test_valid_benzene_connectivity(self): - """Test that a valid benzene conformer has correct connectivity.""" - mol = Chem.MolFromSmiles("c1ccccc1") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - - assert check_connectivity(mol) is True - - def test_valid_cyclohexane_connectivity(self): - """Test that a valid cyclohexane conformer has correct connectivity.""" - mol = Chem.MolFromSmiles("C1CCCCC1") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - - assert check_connectivity(mol) is True - - def test_broken_bond_detected(self): - """Test that a stretched bond is detected as invalid connectivity.""" - mol = Chem.MolFromSmiles("CC") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - - # Manually stretch a C-C bond far apart - conf = mol.GetConformer() - # Move one carbon far away - pos = conf.GetAtomPosition(0) - conf.SetAtomPosition(0, (pos.x + 5.0, pos.y, pos.z)) - - # This should detect the broken bond - assert check_connectivity(mol) is False - - def test_valid_molecule_with_heteroatoms(self): - """Test molecule with nitrogen and oxygen.""" - mol = Chem.MolFromSmiles("CC(=O)NC") # N-methylacetamide - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - - assert check_connectivity(mol) is True - - def test_salt_with_metal_does_not_crash(self): - """Element outside the radii table (Na) must not raise KeyError. - - Sodium acetate contains Na (atomic number 11), which is not in the - UFF radii table. check_connectivity must skip pairs involving an - unknown element rather than indexing the radii dict blindly. - """ - mol = Chem.MolFromSmiles("CC(=O)[O-].[Na+]") # sodium acetate - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - - # Must return a bool without raising KeyError(11). - result = check_connectivity(mol) - assert isinstance(result, bool) - - -class TestModuleImports: - """Every name has exactly one import path: the module that defines it. - - There used to be a second test here importing the same names out of the - ``Auto3D.utils`` package barrel, asserting that both paths worked. The - barrel is gone (``tests/test_import_boundaries.py`` now forbids it), so the - two-paths-for-one-name shape it pinned is the thing being prevented rather - than checked. - """ - - def test_import_from_utils_chemistry(self): - """Test direct import from Auto3D.utils.chemistry.""" - from Auto3D.utils.chemistry import ( - HARTREE_TO_EV, - check_connectivity, - get_mol_charge, - get_rmsd, - hartree2ev, - min_pairwise_distance, - ) - assert HARTREE_TO_EV == hartree2ev - assert callable(get_mol_charge) - assert callable(min_pairwise_distance) - assert callable(get_rmsd) - assert callable(check_connectivity) - - -class TestAmendMol: - """Test the amend_mol function for fixing molecule issues.""" - - def test_amend_mol_preserves_valid_molecule(self): - """Test that a valid molecule is preserved.""" - from Auto3D.utils.chemistry import amend_mol - mol = Chem.MolFromSmiles("CCO") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - - amended_mol = amend_mol(mol) - assert amended_mol is not None - assert amended_mol.GetNumAtoms() == mol.GetNumAtoms() - - def test_amend_mol_returns_none_for_invalid(self): - """Test that amend_mol returns None for severely invalid molecules.""" - from Auto3D.utils.chemistry import amend_mol - # Create molecule and severely distort it - mol = Chem.MolFromSmiles("C") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - - # Severely distort by moving atoms to overlapping positions - conf = mol.GetConformer() - # Move all hydrogens to same position (severe clash) - for i in range(1, mol.GetNumAtoms()): - conf.SetAtomPosition(i, (0.0, 0.0, 0.0)) - - # Should return None for invalid geometry - amended_mol = amend_mol(mol, check_valid=True) - # The function should handle this appropriately - # (either returns None or attempts to fix) - - def test_amend_mol_with_sanitize(self): - """amend_mol(sanitize=True) must actually run RDKit's sanitization, - not just return a non-None object. - - A molecule parsed with ``sanitize=False`` has no ring-perception / - implicit-valence pass yet, so querying ring info raises a - precondition violation -- it genuinely needs sanitizing to become - usable. If ``amend_mol``'s sanitize branch were a no-op, this - molecule would still raise after the call. - """ - from Auto3D.utils.chemistry import amend_mol - - mol = Chem.MolFromSmiles("c1ccccc1", sanitize=False) # benzene, unsanitized - with pytest.raises(RuntimeError): - mol.GetRingInfo().NumRings() # ring perception never ran - - amended_mol = amend_mol(mol, sanitize=True) - - assert amended_mol is not None - # Sanitizing actually ran: ring perception now works and finds the ring. - assert amended_mol.GetRingInfo().NumRings() == 1 - - -class TestGetMolConnectivity: - """Test the get_mol_connectivity function.""" - - def test_ethane_connectivity(self): - """Test connectivity for ethane (C-C single bond). - - Pins the exact canonical ordering (atom1_idx < atom2_idx, per the - function's own docstring/example), not "either order" -- which would - equally accept a broken ``get_mol_connectivity`` that stopped sorting - its tuples. - """ - from Auto3D.utils.chemistry import get_mol_connectivity - mol = Chem.MolFromSmiles("CC") - connectivity = get_mol_connectivity(mol) - - assert connectivity == {(0, 1)} - - def test_ethanol_connectivity(self): - """Test connectivity for ethanol.""" - from Auto3D.utils.chemistry import get_mol_connectivity - mol = Chem.MolFromSmiles("CCO") - connectivity = get_mol_connectivity(mol) - - # Should be a set of tuples - assert isinstance(connectivity, (set, frozenset, list)) - # Should have at least 2 bonds (C-C and C-O) - assert len(connectivity) >= 2 - - def test_benzene_connectivity(self): - """Test connectivity for benzene ring.""" - from Auto3D.utils.chemistry import get_mol_connectivity - mol = Chem.MolFromSmiles("c1ccccc1") - connectivity = get_mol_connectivity(mol) - - # Benzene has 6 C-C bonds in the ring - assert len(connectivity) == 6 - - def test_methane_connectivity(self): - """Test connectivity for methane (no heavy atom bonds).""" - from Auto3D.utils.chemistry import get_mol_connectivity - mol = Chem.MolFromSmiles("C") - connectivity = get_mol_connectivity(mol) - - # Methane has no bonds between heavy atoms (only C-H) - # But if we include H atoms... - mol_with_h = Chem.AddHs(mol) - connectivity_with_h = get_mol_connectivity(mol_with_h) - assert len(connectivity_with_h) == 4 # 4 C-H bonds - - def test_include_bond_order(self): - """Test that bond order can be included. - - The previous version's real assertion sat inside ``if len(bond_info) - == 3:``, which is false exactly when ``include_bond_order`` silently - stops adding the third element -- the one failure mode this test - exists to catch. Assert the 3-tuple shape unconditionally, then the - bond order value. - """ - from Auto3D.utils.chemistry import get_mol_connectivity - mol = Chem.MolFromSmiles("C=C") # Ethene - connectivity = get_mol_connectivity(mol, include_bond_order=True) - - assert connectivity == {(0, 1, 2.0)} - for bond_info in connectivity: - assert len(bond_info) == 3 # (atom1_idx, atom2_idx, bond_order) - assert bond_info[2] == 2.0 # Double bond - - -class TestFilterUnique: - """Test the filter_unique function for RMSD-based duplicate filtering.""" - - def test_filter_identical_conformers(self): - """Test that identical conformers are filtered to one.""" - from Auto3D.utils.chemistry import filter_unique - - mol = Chem.MolFromSmiles("CCO") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - mol.SetProp("Converged", "true") - - # Create identical copies - mol2 = Chem.Mol(mol) - mol2.SetProp("Converged", "true") - - mols = [mol, mol2] - unique_mols = filter_unique(mols, crit=0.3) - - # Should only keep one - assert len(unique_mols) == 1 - - def test_same_geometry_different_energy_kept(self): - """Identical geometry but distinct E_tot must be kept (energy guard). - - Heavy-atom RMSD ~= 0 but the two are distinct minima (the O-H rotamer - case); the energy guard must stop them collapsing into one. - """ - from Auto3D.utils.chemistry import filter_unique - - mol = Chem.MolFromSmiles("CCO") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - mol.SetProp("Converged", "true") - mol.SetProp("E_tot", "-10.0") - - mol2 = Chem.Mol(mol) # identical geometry - mol2.SetProp("Converged", "true") - mol2.SetProp("E_tot", "-10.5") # |dE| >> tol - - unique_mols = filter_unique([mol, mol2], crit=0.3) - assert len(unique_mols) == 2 - - def test_missing_energy_falls_back_to_rmsd_only(self): - """Without E_tot the energy guard cannot apply -> RMSD-only dedup. - - Preserves the legacy behavior for callers that do not set E_tot. - """ - from Auto3D.utils.chemistry import filter_unique - - mol = Chem.MolFromSmiles("CCO") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - mol.SetProp("Converged", "true") # no E_tot set - - mol2 = Chem.Mol(mol) - mol2.SetProp("Converged", "true") - - unique_mols = filter_unique([mol, mol2], crit=0.3) - assert len(unique_mols) == 1 - - def test_filter_different_conformers(self): - """Test that different conformers are kept.""" - from Auto3D.utils.chemistry import filter_unique - - mol1 = Chem.MolFromSmiles("CCCCCC") # Hexane - flexible - mol1 = Chem.AddHs(mol1) - AllChem.EmbedMolecule(mol1, randomSeed=42) - mol1.SetProp("Converged", "true") - - mol2 = Chem.MolFromSmiles("CCCCCC") - mol2 = Chem.AddHs(mol2) - AllChem.EmbedMolecule(mol2, randomSeed=123) - mol2.SetProp("Converged", "true") - - # Generate very different conformers by using different seeds - # and moving atoms around - conf = mol2.GetConformer() - pos = conf.GetAtomPosition(0) - conf.SetAtomPosition(0, (pos.x + 0.5, pos.y, pos.z)) - - mols = [mol1, mol2] - unique_mols = filter_unique(mols, crit=0.3) - - # Should keep both (or at least not crash) - assert len(unique_mols) >= 1 - - def test_two_diastereomers_are_never_merged(self): - """A distinct compound must survive dedup, however close its geometry. - - The same guarantee ``tests/test_filtering.py`` asserts for - ``filter_unique_optimized``, asserted here because this is the other - duplicate filter and it applies the identical RMSD-plus-energy criterion. - Fixing one path and not the other would leave the defect reachable - through ``ConformerRanker(use_optimized_filtering=False)``. - - cis/trans-4-tert-butylcyclohexanol: heavy-atom RMSD between the two - diastereomers was measured at 0.300 A, i.e. at the 0.3 A default - threshold. ``crit`` is opened wide here so RMSD and energy both say - "duplicate" and the only thing that can keep the pair apart is the fact - that they are different compounds. - """ - from Auto3D.utils.chemistry import filter_unique - from Auto3D.utils.energy import set_e_tot_from_ev - - def build(smiles: str) -> Chem.Mol: - mol = Chem.AddHs(Chem.MolFromSmiles(smiles)) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - mol.SetProp("Converged", "true") - set_e_tot_from_ev(mol, -10.0) # identical energies - return mol - - cis = build("O[C@H]1CC[C@@H](CC1)C(C)(C)C") - trans = build("O[C@H]1CC[C@H](CC1)C(C)(C)C") - assert Chem.MolToSmiles(cis) != Chem.MolToSmiles(trans), "test premise" - - assert len(filter_unique([cis, trans], crit=10.0)) == 2, ( - "the legacy filter merged two distinct diastereomers, so an input " - "molecule vanished from the output with no record" - ) - - def test_duplicate_conformers_of_one_stereoisomer_still_collapse(self): - """The other half: the stereo guard must narrow dedup, not disable it.""" - from Auto3D.utils.chemistry import filter_unique - from Auto3D.utils.energy import set_e_tot_from_ev - - def build() -> Chem.Mol: - mol = Chem.AddHs(Chem.MolFromSmiles("O[C@H]1CC[C@@H](CC1)C(C)(C)C")) - AllChem.EmbedMolecule(mol, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol) - mol.SetProp("Converged", "true") - set_e_tot_from_ev(mol, -10.0) - return mol - - assert len(filter_unique([build(), build()], crit=10.0)) == 1, ( - "duplicate conformers of one stereoisomer survived, so the stereo " - "guard has switched dedup off rather than narrowing it" - ) - - def test_filter_unconverged_removed(self): - """Test that unconverged structures are removed.""" - from Auto3D.utils.chemistry import filter_unique - - mol1 = Chem.MolFromSmiles("CCO") - mol1 = Chem.AddHs(mol1) - AllChem.EmbedMolecule(mol1, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol1) - mol1.SetProp("Converged", "true") - - mol2 = Chem.MolFromSmiles("CCO") - mol2 = Chem.AddHs(mol2) - AllChem.EmbedMolecule(mol2, randomSeed=123) - mol2.SetProp("Converged", "false") # Not converged - - mols = [mol1, mol2] - unique_mols = filter_unique(mols, crit=0.3) - - # Only converged one should remain - assert len(unique_mols) == 1 - assert unique_mols[0].GetProp("Converged").lower() == "true" - - def test_filter_empty_list(self): - """Test filtering empty list returns empty list.""" - from Auto3D.utils.chemistry import filter_unique - - unique_mols = filter_unique([], crit=0.3) - assert len(unique_mols) == 0 - - def test_filter_custom_threshold(self): - """Test that custom RMSD threshold works.""" - from Auto3D.utils.chemistry import filter_unique - - mol1 = Chem.MolFromSmiles("CCO") - mol1 = Chem.AddHs(mol1) - AllChem.EmbedMolecule(mol1, randomSeed=42) - AllChem.MMFFOptimizeMolecule(mol1) - mol1.SetProp("Converged", "true") - - mol2 = Chem.Mol(mol1) - mol2.SetProp("Converged", "true") - - mols = [mol1, mol2] - - # With very small threshold, might keep both - unique_mols_small = filter_unique(mols, crit=0.0001) - # With large threshold, definitely keep only one - unique_mols_large = filter_unique(mols, crit=10.0) - - # Large threshold should definitely merge identical mols - assert len(unique_mols_large) == 1 - # A tighter threshold can never merge MORE than a looser one -- the - # discarded half of this test's own computation, now actually checked. - assert len(unique_mols_small) >= len(unique_mols_large) - - def test_filter_unique_removehs_is_linear_and_nondestructive(self, monkeypatch): - """Legacy filter_unique strips Hs once per molecule (not per comparison) and - returns the originals with explicit H + exact positions intact.""" - import numpy as np - from rdkit import Chem - from rdkit.Chem import AllChem - - from Auto3D.utils import chemistry - - base = Chem.AddHs(Chem.MolFromSmiles("CCCCO")) - cids = AllChem.EmbedMultipleConfs(base, numConfs=5, randomSeed=1) - mols = [] - for cid in cids: - m = Chem.Mol(base, confId=int(cid)) - m.SetProp("Converged", "true") - mols.append(m) - n_atoms = base.GetNumAtoms() - orig_pos = {id(m): m.GetConformer().GetPositions().copy() for m in mols} - - calls = {"n": 0} - real_removehs = chemistry.Chem.RemoveHs - - def counting(mol, *a, **k): - calls["n"] += 1 - return real_removehs(mol, *a, **k) - - monkeypatch.setattr(chemistry.Chem, "RemoveHs", counting) - - result = chemistry.filter_unique(mols, crit=0.01) - assert calls["n"] == len(mols) # once per input, never per pair - assert len(result) == len(mols) - for m in result: - assert m.GetNumAtoms() == n_atoms - assert any(a.GetAtomicNum() == 1 for a in m.GetAtoms()) - assert np.array_equal(m.GetConformer().GetPositions(), orig_pos[id(m)]) - - def test_rmsd_failure_keeps_both(self, monkeypatch): - """An incomparable pair (RMSD raises) must NOT be treated as a duplicate. - - When GetBestRMS raises RuntimeError, filter_unique must treat the pair - as distinct (rmsd = inf) and keep both, mirroring the fix already in - filtering._filter_within_cluster. The previous behavior (rmsd = 0) - made distinct conformers look like perfect duplicates and dropped one. - """ - from Auto3D.utils import chemistry - - def make(name): - m = Chem.AddHs(Chem.MolFromSmiles("CCO")) - AllChem.EmbedMolecule(m, randomSeed=abs(hash(name)) % 1000) - AllChem.MMFFOptimizeMolecule(m) - m.SetProp("_Name", name) - m.SetProp("Converged", "true") - return m - - def boom(*args, **kwargs): - raise RuntimeError("GetBestRMS failed") - - # filter_unique calls rdMolAlign.GetBestRMS via the chemistry module. - monkeypatch.setattr(chemistry.rdMolAlign, "GetBestRMS", boom) - - mols = [make("a"), make("b")] - unique_mols = chemistry.filter_unique(mols, crit=0.3) - assert len(unique_mols) == 2 # incomparable pair must NOT be dropped diff --git a/tests/test_utils_connectivity.py b/tests/test_utils_connectivity.py new file mode 100644 index 00000000..3d370b63 --- /dev/null +++ b/tests/test_utils_connectivity.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python +"""Tests for Auto3D.utils.connectivity module.""" +from __future__ import annotations + +import pytest +from rdkit import Chem +from rdkit.Chem import AllChem + +from Auto3D.utils.connectivity import check_connectivity + + +class TestCheckConnectivity: + """Test the check_connectivity function.""" + + def test_valid_ethanol_connectivity(self): + """Test that a valid ethanol conformer has correct connectivity.""" + mol = Chem.MolFromSmiles("CCO") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + + assert check_connectivity(mol) is True + + def test_valid_methane_connectivity(self): + """Test that a valid methane conformer has correct connectivity.""" + mol = Chem.MolFromSmiles("C") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + + assert check_connectivity(mol) is True + + def test_valid_benzene_connectivity(self): + """Test that a valid benzene conformer has correct connectivity.""" + mol = Chem.MolFromSmiles("c1ccccc1") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + + assert check_connectivity(mol) is True + + def test_valid_cyclohexane_connectivity(self): + """Test that a valid cyclohexane conformer has correct connectivity.""" + mol = Chem.MolFromSmiles("C1CCCCC1") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + + assert check_connectivity(mol) is True + + def test_broken_bond_detected(self): + """Test that a stretched bond is detected as invalid connectivity.""" + mol = Chem.MolFromSmiles("CC") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + + # Manually stretch a C-C bond far apart + conf = mol.GetConformer() + # Move one carbon far away + pos = conf.GetAtomPosition(0) + conf.SetAtomPosition(0, (pos.x + 5.0, pos.y, pos.z)) + + # This should detect the broken bond + assert check_connectivity(mol) is False + + def test_valid_molecule_with_heteroatoms(self): + """Test molecule with nitrogen and oxygen.""" + mol = Chem.MolFromSmiles("CC(=O)NC") # N-methylacetamide + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + AllChem.MMFFOptimizeMolecule(mol) + + assert check_connectivity(mol) is True + + def test_salt_with_metal_does_not_crash(self): + """Element outside the radii table (Na) must not raise KeyError. + + Sodium acetate contains Na (atomic number 11), which is not in the + UFF radii table. check_connectivity must skip pairs involving an + unknown element rather than indexing the radii dict blindly. + """ + mol = Chem.MolFromSmiles("CC(=O)[O-].[Na+]") # sodium acetate + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + + # Must return a bool without raising KeyError(11). + result = check_connectivity(mol) + assert isinstance(result, bool) + + +class TestAmendMol: + """Test the amend_mol function for fixing molecule issues.""" + + def test_amend_mol_preserves_valid_molecule(self): + """Test that a valid molecule is preserved.""" + from Auto3D.utils.connectivity import amend_mol + mol = Chem.MolFromSmiles("CCO") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + + amended_mol = amend_mol(mol) + assert amended_mol is not None + assert amended_mol.GetNumAtoms() == mol.GetNumAtoms() + + def test_amend_mol_returns_none_for_invalid(self): + """Test that amend_mol returns None for severely invalid molecules.""" + from Auto3D.utils.connectivity import amend_mol + # Create molecule and severely distort it + mol = Chem.MolFromSmiles("C") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + + # Severely distort by moving atoms to overlapping positions + conf = mol.GetConformer() + # Move all hydrogens to same position (severe clash) + for i in range(1, mol.GetNumAtoms()): + conf.SetAtomPosition(i, (0.0, 0.0, 0.0)) + + # Should return None for invalid geometry + amended_mol = amend_mol(mol, check_valid=True) + # The function should handle this appropriately + # (either returns None or attempts to fix) + + def test_amend_mol_with_sanitize(self): + """amend_mol(sanitize=True) must actually run RDKit's sanitization, + not just return a non-None object. + + A molecule parsed with ``sanitize=False`` has no ring-perception / + implicit-valence pass yet, so querying ring info raises a + precondition violation -- it genuinely needs sanitizing to become + usable. If ``amend_mol``'s sanitize branch were a no-op, this + molecule would still raise after the call. + """ + from Auto3D.utils.connectivity import amend_mol + + mol = Chem.MolFromSmiles("c1ccccc1", sanitize=False) # benzene, unsanitized + with pytest.raises(RuntimeError): + mol.GetRingInfo().NumRings() # ring perception never ran + + amended_mol = amend_mol(mol, sanitize=True) + + assert amended_mol is not None + # Sanitizing actually ran: ring perception now works and finds the ring. + assert amended_mol.GetRingInfo().NumRings() == 1 + + +class TestGetMolConnectivity: + """Test the get_mol_connectivity function.""" + + def test_ethane_connectivity(self): + """Test connectivity for ethane (C-C single bond). + + Pins the exact canonical ordering (atom1_idx < atom2_idx, per the + function's own docstring/example), not "either order" -- which would + equally accept a broken ``get_mol_connectivity`` that stopped sorting + its tuples. + """ + from Auto3D.utils.connectivity import get_mol_connectivity + mol = Chem.MolFromSmiles("CC") + connectivity = get_mol_connectivity(mol) + + assert connectivity == {(0, 1)} + + def test_ethanol_connectivity(self): + """Test connectivity for ethanol.""" + from Auto3D.utils.connectivity import get_mol_connectivity + mol = Chem.MolFromSmiles("CCO") + connectivity = get_mol_connectivity(mol) + + # Should be a set of tuples + assert isinstance(connectivity, (set, frozenset, list)) + # Should have at least 2 bonds (C-C and C-O) + assert len(connectivity) >= 2 + + def test_benzene_connectivity(self): + """Test connectivity for benzene ring.""" + from Auto3D.utils.connectivity import get_mol_connectivity + mol = Chem.MolFromSmiles("c1ccccc1") + connectivity = get_mol_connectivity(mol) + + # Benzene has 6 C-C bonds in the ring + assert len(connectivity) == 6 + + def test_methane_connectivity(self): + """Test connectivity for methane (no heavy atom bonds).""" + from Auto3D.utils.connectivity import get_mol_connectivity + mol = Chem.MolFromSmiles("C") + connectivity = get_mol_connectivity(mol) + + # Methane has no bonds between heavy atoms (only C-H) + # But if we include H atoms... + mol_with_h = Chem.AddHs(mol) + connectivity_with_h = get_mol_connectivity(mol_with_h) + assert len(connectivity_with_h) == 4 # 4 C-H bonds + + def test_include_bond_order(self): + """Test that bond order can be included. + + The previous version's real assertion sat inside ``if len(bond_info) + == 3:``, which is false exactly when ``include_bond_order`` silently + stops adding the third element -- the one failure mode this test + exists to catch. Assert the 3-tuple shape unconditionally, then the + bond order value. + """ + from Auto3D.utils.connectivity import get_mol_connectivity + mol = Chem.MolFromSmiles("C=C") # Ethene + connectivity = get_mol_connectivity(mol, include_bond_order=True) + + assert connectivity == {(0, 1, 2.0)} + for bond_info in connectivity: + assert len(bond_info) == 3 # (atom1_idx, atom2_idx, bond_order) + assert bond_info[2] == 2.0 # Double bond diff --git a/tests/test_utils_energy_constants.py b/tests/test_utils_energy_constants.py new file mode 100644 index 00000000..8c92360d --- /dev/null +++ b/tests/test_utils_energy_constants.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python +"""Tests for the unit-conversion factors and legacy aliases in Auto3D.utils.energy.""" +from __future__ import annotations + +from Auto3D.utils.energy import ( + EV_TO_KCAL_PER_MOL, + HARTREE_TO_EV, + HARTREE_TO_KCAL_PER_MOL, + ev2kcalpermol, + hartree2ev, + hartree2kcalpermol, +) + + +class TestEnergyConversionConstants: + """Test energy conversion constants and aliases.""" + + def test_hartree_to_ev_value(self): + """Test that HARTREE_TO_EV has the correct CODATA 2018 value.""" + assert abs(HARTREE_TO_EV - 27.211386245988) < 1e-10 + + def test_hartree_to_kcal_per_mol_value(self): + """Test that HARTREE_TO_KCAL_PER_MOL has the expected value.""" + assert abs(HARTREE_TO_KCAL_PER_MOL - 627.50947337481) < 1e-8 + + def test_ev_to_kcal_per_mol_value(self): + """Test that EV_TO_KCAL_PER_MOL has the expected value.""" + assert abs(EV_TO_KCAL_PER_MOL - 23.060547830619026) < 1e-10 + + def test_backward_compatibility_aliases(self): + """Test that backward compatibility aliases match constants.""" + assert hartree2ev == HARTREE_TO_EV + assert hartree2kcalpermol == HARTREE_TO_KCAL_PER_MOL + assert ev2kcalpermol == EV_TO_KCAL_PER_MOL + + def test_conversion_consistency(self): + """Test that conversion factors are mathematically consistent.""" + # HARTREE_TO_KCAL_PER_MOL should approximately equal + # HARTREE_TO_EV * EV_TO_KCAL_PER_MOL + calculated = HARTREE_TO_EV * EV_TO_KCAL_PER_MOL + # Allow some tolerance for floating point precision + assert abs(calculated - HARTREE_TO_KCAL_PER_MOL) < 0.001 + + +class TestModuleImports: + """Test that all expected names are importable from their defining modules. + + There used to be a second test here importing the same names out of the + ``Auto3D.utils`` package barrel, asserting that both paths worked. The + barrel is gone (``tests/test_import_boundaries.py`` now forbids it), so the + two-paths-for-one-name shape it pinned is the thing being prevented rather + than checked. + """ + + def test_import_from_defining_modules(self): + """Each name resolves at the module that now defines it.""" + from Auto3D.utils.connectivity import check_connectivity + from Auto3D.utils.energy import HARTREE_TO_EV, hartree2ev + from Auto3D.utils.geometry import get_rmsd, min_pairwise_distance + from Auto3D.utils.molprops import get_mol_charge + + assert HARTREE_TO_EV == hartree2ev + assert callable(get_mol_charge) + assert callable(min_pairwise_distance) + assert callable(get_rmsd) + assert callable(check_connectivity) diff --git a/tests/test_utils_file_ops.py b/tests/test_utils_file_ops.py deleted file mode 100644 index dbfdddc0..00000000 --- a/tests/test_utils_file_ops.py +++ /dev/null @@ -1,1223 +0,0 @@ -"""Tests for Auto3D.utils.file_ops module.""" -from pathlib import Path - -import pytest -from rdkit import Chem - -from Auto3D.utils.file_ops import ( - smiles2smi, - guess_file_type, - hash_enumerated_smi_IDs, - hash_taut_smi, - housekeeping, - create_chunk_meta_names, - combine_smi, - SDF2chunks, - encode_ids, - decode_ids, - reorder_sdf, - count_sdf, - find_ids_not_in_sdf, - find_smiles_not_in_sdf, - iter_smi_records, -) - - -# Get the test files directory -TEST_DIR = Path(__file__).parent -FILES_DIR = TEST_DIR / "files" - - -class TestSmiles2Smi: - """Tests for smiles2smi function.""" - - def test_creates_file_with_inchikeys(self, tmp_path): - """smiles2smi should create a .smi file with SMILES and InChIKey IDs.""" - smiles = ["CCO", "CCC"] - output = tmp_path / "test.smi" - - result = smiles2smi(smiles, str(output)) - - assert result == str(output) - assert output.exists() - content = output.read_text() - lines = content.strip().split('\n') - assert len(lines) == 2 - # Each line should have SMILES and InChIKey - for line in lines: - parts = line.split() - assert len(parts) == 2 - - def test_returns_output_path(self, tmp_path): - """smiles2smi should return the output file path.""" - smiles = ["CCO"] - output = tmp_path / "output.smi" - - result = smiles2smi(smiles, str(output)) - - assert result == str(output) - - def test_inchikey_format(self, tmp_path): - """InChIKeys should have the standard 27-character format.""" - smiles = ["CCO"] - output = tmp_path / "test.smi" - - smiles2smi(smiles, str(output)) - - content = output.read_text().strip() - parts = content.split() - inchikey = parts[1] - # InChIKey format: 14 chars + hyphen + 10 chars + hyphen + 1 char = 27 chars - assert len(inchikey) == 27 - assert inchikey.count('-') == 2 - - def test_preserves_smiles_string(self, tmp_path): - """Original SMILES strings should be preserved in output.""" - smiles = ["C#N", "C=C", "[NH4+]"] - output = tmp_path / "test.smi" - - smiles2smi(smiles, str(output)) - - content = output.read_text() - for smi in smiles: - assert smi in content - - def test_empty_list(self, tmp_path): - """Empty input list should create empty file.""" - output = tmp_path / "test.smi" - - result = smiles2smi([], str(output)) - - assert result == str(output) - assert output.exists() - assert output.read_text() == "" - - def test_colliding_inchikeys_get_distinct_ids(self, tmp_path): - """Two inputs that share an InChIKey must keep distinct IDs. - - The same molecule written two ways (here benzene) yields one InChIKey; - without disambiguation reorder_sdf would collapse the duplicate IDs and - silently drop the second input. Each input must get its own line/ID. - """ - output = tmp_path / "test.smi" - - smiles2smi(["c1ccccc1", "C1=CC=CC=C1"], str(output)) - - lines = output.read_text().strip().split("\n") - assert len(lines) == 2 - ids = [line.split()[1] for line in lines] - assert ids[0] != ids[1], "colliding InChIKeys must be disambiguated" - # First keeps the bare InChIKey; the repeat is suffixed. - assert ids[1] == f"{ids[0]}_2" - - def test_distinct_inputs_keep_bare_inchikeys(self, tmp_path): - """Non-colliding inputs must keep their plain InChIKey IDs (no suffix).""" - output = tmp_path / "test.smi" - - smiles2smi(["CCO", "CCC"], str(output)) - - ids = [line.split()[1] for line in output.read_text().strip().split("\n")] - assert ids[0] != ids[1] - assert all("_" not in i for i in ids) - - -class TestCombineSmiOrderPreservingDedup: - """Tests for combine_smi (order-preserving dedup). - - Named distinctly from the ``TestCombineSmi`` class below -- both defined - ``TestCombineSmi`` until a lint audit found the second definition was - silently shadowing this one, so pytest never collected the test below. - """ - - def test_preserves_order_and_dedups(self, tmp_path): - f1 = tmp_path / "a.smi" - f2 = tmp_path / "b.smi" - f1.write_text("CCO ethanol\nCCC propane\n") - f2.write_text("CCC propane\nCCCC butane\n") # propane duplicated - out = tmp_path / "combined.smi" - - combine_smi([str(f1), str(f2)], str(out)) - - lines = out.read_text().strip().split("\n") - # Deduped (propane once) and in first-seen input order. - assert lines == ["CCO ethanol", "CCC propane", "CCCC butane"] - - -class TestEncodeIdsSmiDenseIndex: - """encode_ids must use a dense, gap-free index for .smi inputs.""" - - def test_blank_lines_do_not_create_index_gaps(self, tmp_path): - smi = tmp_path / "in.smi" - # Blank lines interspersed: the old code used the file line number as the - # index, leaving gaps. The dense record counter must yield 0,1,2. - smi.write_text("CCO a\n\nCCC b\n\n\nCCCC c\n") - - new_path, mapping = encode_ids(str(smi)) - - assert sorted(mapping.values()) == [0, 1, 2] - ids = [line.split()[1] for line in Path(new_path).read_text().strip().split("\n")] - assert ids == ["0", "1", "2"] - - -class TestGuessFileType: - """Tests for guess_file_type function.""" - - def test_smi_extension(self): - """Test detection of .smi files.""" - assert guess_file_type("molecules.smi") == "smi" - assert guess_file_type("/path/to/input.smi") == "smi" - - def test_sdf_extension(self): - """Test detection of .sdf files.""" - assert guess_file_type("molecules.sdf") == "sdf" - assert guess_file_type("/data/output/result.sdf") == "sdf" - - def test_mol2_extension(self): - """Test detection of .mol2 files.""" - assert guess_file_type("molecule.mol2") == "mol2" - - def test_xyz_extension(self): - """Test detection of .xyz files.""" - assert guess_file_type("geometry.xyz") == "xyz" - - def test_complex_path(self): - """Test with complex file paths.""" - assert guess_file_type("/home/user/data.2024/molecules.sdf") == "sdf" - assert guess_file_type("./relative/path/file.smi") == "smi" - - def test_no_extension(self): - """Test file without extension returns empty string.""" - assert guess_file_type("filename") == "" - - def test_hidden_file(self): - """Test hidden files with extension.""" - assert guess_file_type(".hidden.sdf") == "sdf" - - -class TestHashEnumeratedSmiIDs: - """Tests for hash_enumerated_smi_IDs function.""" - - def test_basic_hashing(self, tmp_path): - """Test basic hashing with simple SMILES file.""" - input_file = tmp_path / "input.smi" - output_file = tmp_path / "output.smi" - - # Create input file with unsorted IDs - input_file.write_text("CCO mol_b\nCC mol_a\nCCC mol_c\n") - - hash_enumerated_smi_IDs(str(input_file), str(output_file)) - - # Read and verify output - lines = output_file.read_text().strip().split("\n") - assert len(lines) == 3 - # Should be sorted by ID - assert "mol_a" in lines[0] - assert "mol_b" in lines[1] - assert "mol_c" in lines[2] - - def test_duplicate_id_handling(self, tmp_path): - """Test that duplicate IDs get '_0' suffix.""" - input_file = tmp_path / "input.smi" - output_file = tmp_path / "output.smi" - - # Create input file with duplicate IDs - input_file.write_text("CCO mol1\nCC mol1\nCCC mol1\n") - - hash_enumerated_smi_IDs(str(input_file), str(output_file)) - - lines = output_file.read_text().strip().split("\n") - assert len(lines) == 3 - - # Check that duplicates were renamed - ids = [line.split()[1] for line in lines] - assert "mol1" in ids - assert "mol1_0" in ids - assert "mol1_0_0" in ids - - def test_preserves_smiles(self, tmp_path): - """Test that SMILES strings are preserved correctly.""" - input_file = tmp_path / "input.smi" - output_file = tmp_path / "output.smi" - - input_file.write_text("C#N id1\nC=C id2\n") - - hash_enumerated_smi_IDs(str(input_file), str(output_file)) - - content = output_file.read_text() - assert "C#N" in content - assert "C=C" in content - - -class TestHashTautSmi: - """Tests for hash_taut_smi function.""" - - def test_tautomer_suffix_added(self, tmp_path): - """Test that @taut suffix is added to IDs.""" - input_file = tmp_path / "input.smi" - output_file = tmp_path / "output.smi" - - input_file.write_text("CCO mol1\nCC mol2\n") - - hash_taut_smi(str(input_file), str(output_file)) - - content = output_file.read_text() - assert "@taut" in content - - def test_incremental_taut_suffix(self, tmp_path): - """Test that duplicate base IDs get incrementing taut numbers.""" - input_file = tmp_path / "input.smi" - output_file = tmp_path / "output.smi" - - # Same ID for multiple SMILES - input_file.write_text("CCO mol1\nCC mol1\n") - - hash_taut_smi(str(input_file), str(output_file)) - - lines = output_file.read_text().strip().split("\n") - ids = [line.split()[1] for line in lines] - - # Should have different taut numbers - assert len(set(ids)) == 2 - assert all("@taut" in id for id in ids) - - -class TestHousekeeping: - """Tests for housekeeping function.""" - - def test_moves_files_except_output(self, tmp_path): - """Test that files are moved except for the output file.""" - job_dir = tmp_path / "job" - job_dir.mkdir() - - verbose_folder = tmp_path / "verbose" - verbose_folder.mkdir() - - # Create test files - (job_dir / "meta1.txt").write_text("meta1") - (job_dir / "meta2.txt").write_text("meta2") - output_file = job_dir / "output.sdf" - output_file.write_text("output") - - housekeeping(str(job_dir), str(verbose_folder), str(output_file)) - - # Output should still be in job_dir - assert output_file.exists() - # Meta files should be moved - assert (verbose_folder / "meta1.txt").exists() - assert (verbose_folder / "meta2.txt").exists() - - -class TestCreateChunkMetaNames: - """Tests for create_chunk_meta_names function.""" - - def test_generates_expected_paths(self): - """Test that all expected paths are generated.""" - result = create_chunk_meta_names("chunk1.smi", "/tmp/job") - - assert result["output"] == "/tmp/job/chunk1_3d.sdf" - assert result["optimized_og"] == "/tmp/job/chunk1_3d0.sdf" - assert result["output_taut"] == "/tmp/job/smi_taut.smi" - assert result["smiles_enumerated"] == "/tmp/job/smiles_enumerated.smi" - assert result["smiles_reduced"] == "/tmp/job/smiles_enumerated_reduced.smi" - assert result["smiles_hashed"] == "/tmp/job/smiles_enumerated_hashed.smi" - assert result["enumerated_sdf"] == "/tmp/job/smiles_enumerated.sdf" - assert result["sorted_sdf"] == "/tmp/job/enumerated_sorted.sdf" - assert result["housekeeping_folder"] == "/tmp/job/verbose" - assert result["path"] == "chunk1.smi" - assert result["dir"] == "/tmp/job" - - def test_handles_path_with_directory(self): - """Test that paths with directories work correctly.""" - result = create_chunk_meta_names("/data/input/chunk1.smi", "/output/job") - - assert result["output"] == "/output/job/chunk1_3d.sdf" - assert result["path"] == "/data/input/chunk1.smi" - - -class TestCombineSmi: - """Tests for combine_smi function.""" - - def test_combines_files(self, tmp_path): - """Test that multiple SMILES files are combined.""" - file1 = tmp_path / "file1.smi" - file2 = tmp_path / "file2.smi" - output = tmp_path / "combined.smi" - - file1.write_text("CCO mol1\nCC mol2\n") - file2.write_text("CCC mol3\nCCCC mol4\n") - - combine_smi([str(file1), str(file2)], str(output)) - - content = output.read_text() - assert "mol1" in content - assert "mol2" in content - assert "mol3" in content - assert "mol4" in content - - def test_removes_duplicates(self, tmp_path): - """Test that duplicate entries are removed.""" - file1 = tmp_path / "file1.smi" - file2 = tmp_path / "file2.smi" - output = tmp_path / "combined.smi" - - file1.write_text("CCO mol1\n") - file2.write_text("CCO mol1\n") # Same entry - - combine_smi([str(file1), str(file2)], str(output)) - - lines = output.read_text().strip().split("\n") - assert len(lines) == 1 - - def test_ignores_blank_lines(self, tmp_path): - """Test that blank lines are ignored.""" - file1 = tmp_path / "file1.smi" - output = tmp_path / "combined.smi" - - file1.write_text("CCO mol1\n\n\nCC mol2\n \n") - - combine_smi([str(file1)], str(output)) - - lines = output.read_text().strip().split("\n") - assert len(lines) == 2 - - -class TestSDF2chunks: - """Tests for SDF2chunks function.""" - - def test_splits_sdf_into_chunks(self): - """Test that SDF file is split into molecule chunks.""" - sdf_path = str(FILES_DIR / "example.sdf") - - chunks = SDF2chunks(sdf_path) - - # example.sdf has 2 molecules - assert len(chunks) == 2 - - # Each chunk should end with $$$$ - for chunk in chunks: - assert chunk[-1].strip() == "$$$$" - - def test_chunk_contains_molecule_lines(self): - """Test that chunks contain all molecule lines.""" - sdf_path = str(FILES_DIR / "example.sdf") - - chunks = SDF2chunks(sdf_path) - - # First chunk should start with molecule name - assert chunks[0][0].strip() == "mol1" - assert chunks[1][0].strip() == "mol2" - - def test_preserves_all_content(self): - """Test that all content from original file is preserved.""" - sdf_path = str(FILES_DIR / "example.sdf") - - chunks = SDF2chunks(sdf_path) - - # Reconstruct file from chunks - reconstructed = "".join(line for chunk in chunks for line in chunk) - - with open(sdf_path) as f: - original = f.read() - - assert reconstructed == original - - -class TestEncodeDecodeIds: - """Tests for encode_ids and decode_ids functions.""" - - def test_encode_smi_file(self, tmp_path): - """Test encoding IDs in a SMILES file.""" - input_file = tmp_path / "input.smi" - input_file.write_text("CCO mol_alpha\nCC mol_beta\nCCC mol_gamma\n") - - new_path, mapping = encode_ids(str(input_file)) - - assert mapping == {"mol_alpha": 0, "mol_beta": 1, "mol_gamma": 2} - assert Path(new_path).name == "input_encoded.smi" - - # Check encoded file content - content = Path(new_path).read_text() - assert "CCO 0" in content - assert "CC 1" in content - assert "CCC 2" in content - - def test_encode_sdf_file(self): - """Test encoding IDs in an SDF file.""" - sdf_path = str(FILES_DIR / "example.sdf") - - new_path, mapping = encode_ids(sdf_path) - - assert "mol1" in mapping - assert "mol2" in mapping - assert mapping["mol1"] == 0 - assert mapping["mol2"] == 1 - - # Clean up - Path(new_path).unlink(missing_ok=True) - - def test_encode_invalid_extension_raises(self, tmp_path): - """Test that invalid file extension raises ValueError.""" - input_file = tmp_path / "input.xyz" - input_file.write_text("invalid") - - with pytest.raises(ValueError, match="smi or sdf"): - encode_ids(str(input_file)) - - def test_encode_skips_blank_lines(self, tmp_path): - """Test that blank lines in SMILES file are skipped.""" - input_file = tmp_path / "input.smi" - input_file.write_text("CCO mol1\n\n \nCC mol2\n") - - new_path, mapping = encode_ids(str(input_file)) - - # mapping indices may not be sequential if blank lines are in between - assert len(mapping) == 2 - - # Clean up - Path(new_path).unlink(missing_ok=True) - - def test_encode_ids_rejects_duplicate_ids(self, tmp_path): - """Duplicate molecule IDs in a .smi file are rejected up front.""" - from Auto3D.exceptions import InputValidationError - - p = tmp_path / "dup.smi" - p.write_text("CCO mol1\nCCC mol1\n") - with pytest.raises(InputValidationError, match="[Dd]uplicate"): - encode_ids(str(p)) - - def test_encode_ids_rejects_missing_id(self, tmp_path): - """A .smi row without a whitespace-separated ID is rejected.""" - from Auto3D.exceptions import InputValidationError - - p = tmp_path / "noid.smi" - p.write_text("CCO\n") # no whitespace-separated ID - with pytest.raises(InputValidationError, match="ID"): - encode_ids(str(p)) - - def test_encode_ids_roundtrip_unique(self, tmp_path): - """Unique IDs encode cleanly and appear in the mapping.""" - p = tmp_path / "ok.smi" - p.write_text("CCO a\nCCC b\n") - _, mapping = encode_ids(str(p)) - assert set(mapping) == {"a", "b"} - - def test_encode_ids_rejects_blank_sdf_name(self, tmp_path): - """A molecule with a blank _Name in a .sdf file is rejected.""" - from rdkit import Chem - from rdkit.Chem import AllChem - - from Auto3D.exceptions import InputValidationError - - sdf = tmp_path / "blank.sdf" - with Chem.SDWriter(str(sdf)) as w: - m = Chem.AddHs(Chem.MolFromSmiles("CCO")) - AllChem.EmbedMolecule(m, randomSeed=1) - m.SetProp("_Name", "") # blank name - w.write(m) - with pytest.raises(InputValidationError): - encode_ids(str(sdf)) - - def test_encode_ids_refuses_to_overwrite_an_existing_file(self, tmp_path): - """The `_encoded.` name belongs to the user until proven - otherwise. - - The name is derived from the input, so `mols_encoded.smi` beside - `mols.smi` is an ordinary thing for a user to own -- and this function - used to open it for writing without a word. `WorkflowOrchestrator` - now redirects the encoded copy into its own job directory (see - `out_dir` below), but this check keeps the guarantee attached to the - function itself, so a caller taking the default location cannot - reintroduce the defect. - """ - from Auto3D.exceptions import ConfigurationError - - p = tmp_path / "mols.smi" - p.write_text("CCO a\n") - users_file = tmp_path / "mols_encoded.smi" - users_file.write_bytes(b"IRREPLACEABLE USER DATA\n") - - with pytest.raises(ConfigurationError, match="would overwrite"): - encode_ids(str(p)) - - assert users_file.read_bytes() == b"IRREPLACEABLE USER DATA\n" - - def test_encode_ids_writes_into_out_dir_when_given_one(self, tmp_path): - """`out_dir` moves the encoded copy somewhere the caller owns. - - This is how the run pipeline avoids the collision above entirely: it - passes the job directory it just created. The file name is unchanged, - only its directory -- downstream code (`_setup_job_directory`, - `decode_ids`) parses that name. - """ - p = tmp_path / "mols.smi" - p.write_text("CCO a\nCCC b\n") - staging = tmp_path / "staging" - staging.mkdir() - - new_path, mapping = encode_ids(str(p), out_dir=staging) - - assert Path(new_path).parent == staging - assert Path(new_path).name == "mols_encoded.smi" - assert mapping == {"a": 0, "b": 1} - assert not (tmp_path / "mols_encoded.smi").exists() - - -class TestReorderSdf: - """Tests for reorder_sdf function.""" - - def test_reorder_sdf_from_smi(self, tmp_path): - """Test reordering SDF file based on SMILES file order.""" - from rdkit import Chem - - # Create source SMILES file with specific order - smi_file = tmp_path / "source.smi" - smi_file.write_text("CCO mol_b\nCC mol_a\nCCC mol_c\n") - - # Create SDF file with different order - sdf_file = tmp_path / "mols.sdf" - writer = Chem.SDWriter(str(sdf_file)) - for name in ["mol_a", "mol_c", "mol_b"]: - mol = Chem.MolFromSmiles("C") - mol.SetProp("_Name", name) - writer.write(mol) - writer.close() - - # Reorder - result = reorder_sdf(str(sdf_file), str(smi_file)) - - # Verify order matches source - assert len(result) == 3 - assert result[0].GetProp("_Name") == "mol_b" - assert result[1].GetProp("_Name") == "mol_a" - assert result[2].GetProp("_Name") == "mol_c" - - def test_reorder_sdf_from_sdf(self, tmp_path): - """Test reordering SDF file based on another SDF file order.""" - from rdkit import Chem - - # Create source SDF file with specific order - source_sdf = tmp_path / "source.sdf" - writer = Chem.SDWriter(str(source_sdf)) - for name in ["mol_x", "mol_y", "mol_z"]: - mol = Chem.MolFromSmiles("C") - mol.SetProp("_Name", name) - writer.write(mol) - writer.close() - - # Create target SDF file with different order - target_sdf = tmp_path / "target.sdf" - writer = Chem.SDWriter(str(target_sdf)) - for name in ["mol_z", "mol_x", "mol_y"]: - mol = Chem.MolFromSmiles("C") - mol.SetProp("_Name", name) - writer.write(mol) - writer.close() - - # Reorder - result = reorder_sdf(str(target_sdf), str(source_sdf)) - - # Verify order matches source - assert len(result) == 3 - assert result[0].GetProp("_Name") == "mol_x" - assert result[1].GetProp("_Name") == "mol_y" - assert result[2].GetProp("_Name") == "mol_z" - - def test_reorder_sdf_with_tautomers(self, tmp_path): - """Test reordering handles tautomer IDs correctly.""" - from rdkit import Chem - - # Create source SMILES file - smi_file = tmp_path / "source.smi" - smi_file.write_text("CCO mol1\nCC mol2\n") - - # Create SDF file with tautomer variants - sdf_file = tmp_path / "mols.sdf" - writer = Chem.SDWriter(str(sdf_file)) - for name in ["mol2@taut1", "mol1@taut1", "mol1@taut2"]: - mol = Chem.MolFromSmiles("C") - mol.SetProp("_Name", name) - writer.write(mol) - writer.close() - - # Reorder - result = reorder_sdf(str(sdf_file), str(smi_file)) - - # Verify mol1 variants come before mol2 variants - assert len(result) == 3 - # mol1 should be first (2 tautomers) - assert "mol1" in result[0].GetProp("_Name") - assert "mol1" in result[1].GetProp("_Name") - # mol2 should be last - assert "mol2" in result[2].GetProp("_Name") - - def test_reorder_sdf_unsupported_format(self, tmp_path, caplog): - """Test that unsupported format returns None.""" - import logging - - xyz_file = tmp_path / "source.xyz" - xyz_file.write_text("invalid") - - sdf_file = tmp_path / "mols.sdf" - sdf_file.write_text("dummy") - - with caplog.at_level(logging.WARNING): - result = reorder_sdf(str(sdf_file), str(xyz_file)) - - assert result is None - assert "Unsupported file format" in caplog.text - - -class TestFileOpsIntegration: - """Integration tests for file_ops module.""" - - def test_create_chunks_and_housekeeping_workflow(self, tmp_path): - """Test a typical workflow using multiple file_ops functions.""" - # Create job directory structure - job_dir = tmp_path / "job" - job_dir.mkdir() - - # Create meta names - meta = create_chunk_meta_names("input.smi", str(job_dir)) - - # Verify structure - assert "verbose" in meta["housekeeping_folder"] - - # Create verbose folder - Path(meta["housekeeping_folder"]).mkdir() - - # Create some intermediate files - Path(meta["smiles_enumerated"]).write_text("CCO mol1\n") - Path(meta["output"]).write_text("fake sdf output") - - # Run housekeeping - should move enumerated but not output - housekeeping( - str(job_dir), - meta["housekeeping_folder"], - meta["output"] - ) - - # Output should still exist - assert Path(meta["output"]).exists() - # Enumerated should be moved to verbose folder - assert (Path(meta["housekeeping_folder"]) / "smiles_enumerated.smi").exists() - - -def _make_mol(name): - """Build a tiny named RDKit mol for SDF round-trips.""" - from rdkit import Chem - - mol = Chem.MolFromSmiles("C") - mol.SetProp("_Name", name) - return mol - - -class TestNoneMolHardening: - """FIX 1: None records yielded by SDMolSupplier must not crash these helpers. - - A single unparseable SDF record makes SDMolSupplier yield ``None``. The - iterating helpers previously called ``mol.GetProp(...)`` / ``mol.GetNumAtoms()`` - on it and raised ``AttributeError``. They must skip ``None`` instead. - """ - - def test_count_sdf_skips_none_records(self, tmp_path, monkeypatch): - """count_sdf must not count (or crash on) a None record.""" - - import Auto3D.utils.file_ops as file_ops - - valid = _make_mol("mol_a") - monkeypatch.setattr( - file_ops.Chem, "SDMolSupplier", lambda *a, **k: [valid, None] - ) - - sdf = tmp_path / "mols.sdf" - sdf.write_text("placeholder") # path only needs to exist for the call - - assert count_sdf(str(sdf)) == 1 - - def test_decode_ids_skips_none_records(self, tmp_path, monkeypatch): - """decode_ids must skip None records without raising.""" - from rdkit import Chem - - import Auto3D.utils.file_ops as file_ops - - valid = Chem.MolFromSmiles("C") - valid.SetProp("_Name", "0") - valid.SetProp("ID", "0_conf1") - - monkeypatch.setattr( - file_ops.Chem, "SDMolSupplier", lambda *a, **k: [valid, None] - ) - - # decode_ids expects a stem with at least two underscore parts. - sdf = tmp_path / "mols_3d_encoded.sdf" - sdf.write_text("placeholder") - - out = decode_ids(str(sdf), {"mol_a": 0}) - # Only the valid record is written; no AttributeError on the None. - written = count_sdf(out) - assert written == 1 - - def test_find_smiles_not_in_sdf_skips_none_records(self, tmp_path, monkeypatch): - """find_smiles_not_in_sdf must skip None SDF records.""" - from rdkit import Chem - - import Auto3D.utils.file_ops as file_ops - - valid = Chem.MolFromSmiles("C") - valid.SetProp("_Name", "mol_a") - monkeypatch.setattr( - file_ops.Chem, "SDMolSupplier", lambda *a, **k: [valid, None] - ) - - smi = tmp_path / "in.smi" - smi.write_text("C mol_a\nCC mol_b\n") - sdf = tmp_path / "out.sdf" - sdf.write_text("placeholder") - - bad = find_smiles_not_in_sdf(str(smi), str(sdf)) - # mol_a is present (valid mol), mol_b is missing -> reported. - assert ("mol_b", "CC") in bad - assert all(mol_id != "mol_a" for mol_id, _ in bad) - - def test_reorder_sdf_skips_none_records(self, tmp_path, monkeypatch): - """reorder_sdf must skip None records in the target SDF.""" - - import Auto3D.utils.file_ops as file_ops - - smi = tmp_path / "source.smi" - smi.write_text("C mol_a\nC mol_b\n") - - valid_a = _make_mol("mol_a") - valid_b = _make_mol("mol_b") - monkeypatch.setattr( - file_ops.Chem, "SDMolSupplier", lambda *a, **k: [valid_a, None, valid_b] - ) - - sdf = tmp_path / "target.sdf" - sdf.write_text("placeholder") - - result = reorder_sdf(str(sdf), str(smi)) - names = [m.GetProp("_Name") for m in result] - assert names == ["mol_a", "mol_b"] - - -class TestReorderSdfDataPreservation: - """FIX 2: reorder_sdf must not drop unmatched molecules or truncate input.""" - - def test_unmatched_mol_is_preserved(self, tmp_path): - """A mol whose id is not in the source must still survive to disk.""" - from rdkit import Chem - - smi = tmp_path / "source.smi" - # source lists only mol_a and mol_b; mol_c is unmatched. - smi.write_text("C mol_a\nC mol_b\n") - - sdf = tmp_path / "target.sdf" - writer = Chem.SDWriter(str(sdf)) - for name in ["mol_c", "mol_b", "mol_a"]: - writer.write(_make_mol(name)) - writer.close() - - result = reorder_sdf(str(sdf), str(smi)) - - # All three mols preserved (no silent data loss). - result_names = [m.GetProp("_Name") for m in result] - assert set(result_names) == {"mol_a", "mol_b", "mol_c"} - assert len(result_names) == 3 - # Matched ids appear first, in source order. - assert result_names[0] == "mol_a" - assert result_names[1] == "mol_b" - - # And the on-disk file must contain all three as well. - on_disk = [m.GetProp("_Name") for m in Chem.SDMolSupplier(str(sdf))] - assert set(on_disk) == {"mol_a", "mol_b", "mol_c"} - - def test_normal_all_matched_ordering_unchanged(self, tmp_path): - """When every id is matched, ordering is exactly the source order.""" - from rdkit import Chem - - smi = tmp_path / "source.smi" - smi.write_text("C mol_b\nC mol_a\nC mol_c\n") - - sdf = tmp_path / "target.sdf" - writer = Chem.SDWriter(str(sdf)) - for name in ["mol_a", "mol_c", "mol_b"]: - writer.write(_make_mol(name)) - writer.close() - - result = reorder_sdf(str(sdf), str(smi)) - names = [m.GetProp("_Name") for m in result] - assert names == ["mol_b", "mol_a", "mol_c"] - - -class TestSDF2chunksTrailingRecord: - """FIX 3: a final record lacking the $$$$ terminator must not be dropped.""" - - def test_trailing_record_without_terminator_preserved(self, tmp_path): - """SDF2chunks keeps a terminator-less trailing record as its own chunk.""" - sdf = tmp_path / "ragged.sdf" - # First record has $$$$; second record lacks it. - sdf.write_text( - "mol1\n line1\n$$$$\n" - "mol2\n line2\n line3\n" - ) - - chunks = SDF2chunks(str(sdf)) - - assert len(chunks) == 2 - assert chunks[0][0].strip() == "mol1" - # The trailing record's lines must be present in the final chunk. - assert chunks[1][0].strip() == "mol2" - joined = "".join(chunks[1]) - assert "line2" in joined - assert "line3" in joined - - -class TestSmiles2SmiInvalidInput: - """FIX 4: smiles2smi must raise a clear error on an invalid SMILES.""" - - def test_invalid_smiles_raises_input_validation_error(self, tmp_path): - """An unparseable SMILES raises InputValidationError naming the SMILES.""" - from Auto3D.exceptions import InputValidationError - - out = tmp_path / "out.smi" - with pytest.raises(InputValidationError, match=r"C\(C"): - smiles2smi(["CCO", "C(C"], str(out)) - - -class TestHashHelpersBlankLines: - """FIX 5: blank / malformed lines must not crash the hashing helpers.""" - - def test_hash_enumerated_skips_blank_and_extra_token_lines(self, tmp_path): - """hash_enumerated_smi_IDs tolerates blank lines and extra tokens.""" - inp = tmp_path / "in.smi" - inp.write_text("CCO mol1\n\n \nCC mol2 extra_token\n") - out = tmp_path / "out.smi" - - # Must not raise ValueError. - hash_enumerated_smi_IDs(str(inp), str(out)) - - lines = [ln for ln in out.read_text().splitlines() if ln.strip()] - ids = [ln.split()[1] for ln in lines] - assert "mol1" in ids - assert "mol2" in ids - - def test_hash_taut_skips_blank_and_extra_token_lines(self, tmp_path): - """hash_taut_smi tolerates blank lines and extra tokens.""" - inp = tmp_path / "in.smi" - inp.write_text("CCO mol1\n\nCC mol2 extra_token\n") - out = tmp_path / "out.smi" - - hash_taut_smi(str(inp), str(out)) - - lines = [ln for ln in out.read_text().splitlines() if ln.strip()] - assert len(lines) == 2 - assert all("@taut" in ln.split()[1] for ln in lines) - - def test_find_smiles_not_in_sdf_tolerates_blank_and_3token_lines( - self, tmp_path, monkeypatch - ): - """find_smiles_not_in_sdf tolerates blank and 3-token .smi lines.""" - from rdkit import Chem - - import Auto3D.utils.file_ops as file_ops - - valid = Chem.MolFromSmiles("C") - valid.SetProp("_Name", "mol_a") - monkeypatch.setattr( - file_ops.Chem, "SDMolSupplier", lambda *a, **k: [valid] - ) - - smi = tmp_path / "in.smi" - # blank line + a 3-token line (first two tokens taken). - smi.write_text("C mol_a\n\nCC mol_b extra\n") - sdf = tmp_path / "out.sdf" - sdf.write_text("placeholder") - - bad = find_smiles_not_in_sdf(str(smi), str(sdf)) - assert ("mol_b", "CC") in bad - - -class TestIterSmiRecords: - """FIX A: shared lenient .smi parser used by all 7 call sites.""" - - def test_blank_lines_skipped(self, tmp_path): - """Blank and whitespace-only lines yield no records.""" - p = tmp_path / "in.smi" - p.write_text("CCO mol1\n\n \nCC mol2\n") - records = list(iter_smi_records(str(p))) - assert [(s, i) for _ln, s, i in records] == [("CCO", "mol1"), ("CC", "mol2")] - # line_no is 1-based and reflects the original line position. - assert records[0][0] == 1 - assert records[1][0] == 4 - - def test_three_token_line_yields_first_two(self, tmp_path): - """A 3-token line yields only the first two tokens (extras ignored).""" - p = tmp_path / "in.smi" - p.write_text("CCN extra_a extra_b\n") - records = list(iter_smi_records(str(p))) - assert len(records) == 1 - line_no, smiles, mol_id = records[0] - assert (smiles, mol_id) == ("CCN", "extra_a") - - def test_on_malformed_skip_skips_one_token_line_with_warning( - self, tmp_path, caplog - ): - """on_malformed='skip' (default) skips a 1-token line and warns.""" - import logging - - p = tmp_path / "in.smi" - p.write_text("CCO mol1\nC1CCCCC1\nCC mol2\n") - with caplog.at_level(logging.WARNING): - records = list(iter_smi_records(str(p), on_malformed="skip")) - assert [(s, i) for _ln, s, i in records] == [("CCO", "mol1"), ("CC", "mol2")] - assert any("failed to parse" in r.message for r in caplog.records) - - def test_on_malformed_raise_raises_on_one_token_line(self, tmp_path): - """on_malformed='raise' raises InputValidationError naming the line.""" - from Auto3D.exceptions import InputValidationError - - p = tmp_path / "in.smi" - p.write_text("CCO mol1\nC1CCCCC1\n") - with pytest.raises(InputValidationError, match="Line 2"): - list(iter_smi_records(str(p), on_malformed="raise")) - - def test_invalid_on_malformed_value_raises(self, tmp_path): - """An unknown on_malformed value raises ValueError.""" - p = tmp_path / "in.smi" - p.write_text("CCO mol1\n") - with pytest.raises(ValueError, match="on_malformed"): - list(iter_smi_records(str(p), on_malformed="bogus")) - - -def test_housekeeping_sweep_is_per_file_robust(tmp_path, monkeypatch): - """One unmovable file must not abandon the rest of the sweep. - - This guard used to live on a second loop that swept `oeomega_*` out of the - *process working directory*; that loop is gone (it destroyed user files -- - see `TestHousekeepingStaysInsideTheJobDirectory` in tests/test_durability.py) - and the OpenEye logfiles it collected now land inside the job directory, - where this loop picks them up. The robustness property moved with them: a - permission error, or a file that vanished under us, must leave a complete - `verbose` folder minus that one file rather than a half-populated one plus - a traceback out of `optim_rank_wrapper`'s blanket except. - """ - import os - - from Auto3D.utils.file_ops import housekeeping - - job = tmp_path / "job" - job.mkdir() - dest = tmp_path / "verbose" - dest.mkdir() - - # Two logfiles in the job directory; the FIRST one encountered (by - # counter) will fail to move. - (job / "oeomega_a.log").write_text("a") - (job / "oeomega_b.log").write_text("b") - - real_move = __import__("shutil").move - call_count = {"n": 0} - - def flaky_move(src, dst): - call_count["n"] += 1 - if call_count["n"] == 1: - # Simulate the file having gone away underneath the sweep. - if os.path.exists(src): - os.remove(src) - raise OSError("already gone") - return real_move(src, dst) - - monkeypatch.setattr("Auto3D.utils.file_ops.shutil.move", flaky_move) - - housekeeping(str(job), str(dest), str(job / "out.sdf")) # must not raise - - # Exactly one of the two logfiles must have been successfully moved: the - # sweep continued past the failure instead of stopping on it. - moved = list(dest.glob("oeomega_*.log")) - assert len(moved) == 1, f"Expected 1 moved file, got {[f.name for f in moved]}" - - -class TestFindSmilesNotInSdfTautStripping: - """C7: a decoded '@tautN' suffix must not cause a false "missing" report.""" - - def test_taut_suffixed_output_name_matches_base_smi_id(self, tmp_path): - """decode_ids keeps 'id@tautN' on tautomer conformers; the .smi only - has the base id, so find_smiles_not_in_sdf must strip the suffix - before comparing or every tautomer-derived molecule is misreported.""" - smi = tmp_path / "in.smi" - smi.write_text("CCO mol_a\n") - - sdf = tmp_path / "out.sdf" - writer = Chem.SDWriter(str(sdf)) - writer.write(_make_mol("mol_a@taut0")) - writer.close() - - bad = find_smiles_not_in_sdf(str(smi), str(sdf)) - assert bad == [], f"mol_a wrongly reported missing: {bad}" - - -class TestFindIdsNotInSdf: - """find_ids_not_in_sdf: the SDF-input counterpart to find_smiles_not_in_sdf.""" - - def test_missing_id_is_reported(self, tmp_path): - """An id present in the source SDF but absent from the output SDF is reported.""" - source = tmp_path / "source.sdf" - writer = Chem.SDWriter(str(source)) - for name in ["mol_a", "mol_b"]: - writer.write(_make_mol(name)) - writer.close() - - out = tmp_path / "out.sdf" - writer = Chem.SDWriter(str(out)) - writer.write(_make_mol("mol_a")) # mol_b never produced a structure - writer.close() - - bad = find_ids_not_in_sdf(str(source), str(out)) - assert bad == ["mol_b"] - - def test_no_missing_ids_returns_empty_list(self, tmp_path): - """Every source id present in the output -> nothing reported.""" - source = tmp_path / "source.sdf" - writer = Chem.SDWriter(str(source)) - for name in ["mol_a", "mol_b"]: - writer.write(_make_mol(name)) - writer.close() - - out = tmp_path / "out.sdf" - writer = Chem.SDWriter(str(out)) - for name in ["mol_b", "mol_a"]: - writer.write(_make_mol(name)) - writer.close() - - assert find_ids_not_in_sdf(str(source), str(out)) == [] - - def test_taut_suffixed_output_name_matches_base_id(self, tmp_path): - """Same '@tautN' stripping as find_smiles_not_in_sdf, for SDF input.""" - source = tmp_path / "source.sdf" - writer = Chem.SDWriter(str(source)) - writer.write(_make_mol("mol_a")) - writer.close() - - out = tmp_path / "out.sdf" - writer = Chem.SDWriter(str(out)) - writer.write(_make_mol("mol_a@taut1")) - writer.close() - - assert find_ids_not_in_sdf(str(source), str(out)) == [] - - def test_an_unreadable_record_is_reported_on_the_input_side_only( - self, tmp_path, monkeypatch - ): - """The two sides of the comparison are not symmetric, and must not be. - - This test used to assert ``== []`` for a source file containing an - unreadable record, under the heading "must not crash or miscount". The - empty list *was* the miscount: an input molecule the pipeline never saw - was reported by nothing, and the run exited 0 claiming completeness. - - The asymmetry is the point: - - * **input side** -- a record that cannot be read is a molecule the user - supplied and did not get back. It is reported, by position, since it has - no ``_Name`` to report by. - * **output side** -- a record that cannot be read yields no name to match - against, and there is nothing better to do than skip it. Reporting it - would invent a *missing input* out of an unreadable output. - """ - import Auto3D.utils.file_ops as file_ops - - calls = {"n": 0} - - def fake_supplier(*a, **k): - calls["n"] += 1 - if calls["n"] == 1: - return [_make_mol("mol_a"), None] # source: one good, one unreadable - return [None, _make_mol("mol_a")] # output: mol_a made it - - monkeypatch.setattr(file_ops.Chem, "SDMolSupplier", fake_supplier) - - source = tmp_path / "source.sdf" - source.write_text("placeholder") - out = tmp_path / "out.sdf" - out.write_text("placeholder") - - assert find_ids_not_in_sdf(str(source), str(out)) == [ - file_ops.UNPARSEABLE_RECORD_ID.format(index=1) - ] - - -class TestReconciliationSeesUnparseableInputRecords: - """A record Auto3D could not read must not vanish from the accounting. - - ``encode_ids`` skips an unparseable SDF record with a warning, so it never - enters the run. ``find_ids_not_in_sdf`` then built its expected-ID list by - reading **the same source SDF** and skipping the same record -- so it was in - neither ``source_ids`` nor the output, could not appear in ``failures``, and - ``_exit_if_incomplete`` saw ``failed_count == 0``. The run printed a success - summary and exited **0** having processed fewer molecules than the file - contained, which is precisely what the C7 reconciliation exists to prevent. - - Only the SDF path is affected. ``encode_ids`` reads ``.smi`` input with - ``on_malformed="raise"``, so a malformed SMILES line aborts the run with - ``InputValidationError`` long before reconciliation -- the same blindness - cannot be reached through that door. - """ - - @staticmethod - def _unparseable_record(name: str) -> str: - """A molblock RDKit rejects: the counts line is corrupted. - - Built from a real molblock so only the one line under test is invalid; - a wholly invented block could fail for an unrelated reason. - """ - mol = _make_mol(name) - lines = Chem.MolToMolBlock(mol).splitlines() - lines[3] = "!! corrupted counts line !!" - return "\n".join(lines) - - def test_an_unparseable_source_record_is_reported_as_a_failure(self, tmp_path): - source = tmp_path / "source.sdf" - good = Chem.MolToMolBlock(_make_mol("mol_a")) - source.write_text( - good + "$$$$\n" + self._unparseable_record("mol_b") + "\n$$$$\n" - ) - # Confirm the premise: RDKit reads one molecule and one None. - parsed = list(Chem.SDMolSupplier(str(source), removeHs=False)) - assert [m is None for m in parsed] == [False, True], "test premise" - - out = tmp_path / "out.sdf" - writer = Chem.SDWriter(str(out)) - writer.write(_make_mol("mol_a")) # the parseable one succeeded - writer.close() - - missing = find_ids_not_in_sdf(str(source), str(out)) - - assert len(missing) == 1, ( - f"a source record the pipeline could not read was left out of the " - f"accounting entirely, so the run would exit 0 claiming every input " - f"was processed; got {missing}" - ) - assert "1" in missing[0], ( - f"the report must say which record could not be read, got {missing[0]!r}" - ) - - def test_a_clean_source_file_still_reports_nothing(self, tmp_path): - """The new branch must not manufacture failures for a healthy file.""" - source = tmp_path / "source.sdf" - writer = Chem.SDWriter(str(source)) - for name in ("mol_a", "mol_b"): - writer.write(_make_mol(name)) - writer.close() - - out = tmp_path / "out.sdf" - writer = Chem.SDWriter(str(out)) - for name in ("mol_a", "mol_b"): - writer.write(_make_mol(name)) - writer.close() - - assert find_ids_not_in_sdf(str(source), str(out)) == [] - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/test_utils_geometry.py b/tests/test_utils_geometry.py new file mode 100644 index 00000000..e8856cd7 --- /dev/null +++ b/tests/test_utils_geometry.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python +"""Tests for Auto3D.utils.geometry module.""" +from __future__ import annotations + +import numpy as np +import pytest # noqa: F401 (several tests below are parametrized helpers' home) +from rdkit import Chem +from rdkit.Chem import AllChem + +from Auto3D.utils.geometry import get_rmsd, min_pairwise_distance + + +class TestMinPairwiseDistance: + """Test the min_pairwise_distance function.""" + + def test_simple_three_points(self): + """Test with three simple points.""" + points = np.array([ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 2.0, 0.0] + ]) + result = min_pairwise_distance(points) + assert abs(result - 1.0) < 1e-5 + + def test_two_points(self): + """Test with two points.""" + points = np.array([ + [0.0, 0.0, 0.0], + [3.0, 4.0, 0.0] # Distance = 5 + ]) + result = min_pairwise_distance(points) + assert abs(result - 5.0) < 1e-5 + + def test_collinear_points(self): + """Test with collinear points.""" + points = np.array([ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [2.0, 0.0, 0.0], + [5.0, 0.0, 0.0] + ]) + result = min_pairwise_distance(points) + assert abs(result - 1.0) < 1e-5 + + def test_3d_points(self): + """Test with points in 3D space.""" + points = np.array([ + [0.0, 0.0, 0.0], + [1.0, 1.0, 1.0], # Distance = sqrt(3) ~ 1.732 + [5.0, 5.0, 5.0] + ]) + result = min_pairwise_distance(points) + expected = np.sqrt(3) + assert abs(result - expected) < 1e-5 + + def test_input_type_conversion(self): + """Test that integer input is properly converted.""" + points = np.array([ + [0, 0, 0], + [1, 0, 0], + [0, 2, 0] + ], dtype=np.int32) + result = min_pairwise_distance(points) + assert abs(result - 1.0) < 1e-5 + + def test_very_close_points(self): + """Test with very close points.""" + points = np.array([ + [0.0, 0.0, 0.0], + [0.001, 0.0, 0.0], + [10.0, 0.0, 0.0] + ]) + result = min_pairwise_distance(points) + assert abs(result - 0.001) < 1e-6 + + +class TestGetRmsd: + """Test the get_rmsd function.""" + + def test_identical_molecules(self): + """Test RMSD of a molecule with itself is 0.""" + mol = Chem.MolFromSmiles("CCO") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + + mol_copy = Chem.Mol(mol) + rmsd = get_rmsd(mol, mol_copy) + assert abs(rmsd) < 1e-5 + + def test_different_conformers(self): + """Test RMSD of different conformers is > 0.""" + mol = Chem.MolFromSmiles("CCCCCC") # Hexane - flexible + mol = Chem.AddHs(mol) + + # Generate two different conformers + AllChem.EmbedMolecule(mol, randomSeed=42) + conf1 = mol.GetConformer() + + mol2 = Chem.Mol(mol) + AllChem.EmbedMolecule(mol2, randomSeed=123) + + rmsd = get_rmsd(mol, mol2) + # Different random seeds should give different conformers + # The RMSD should be >= 0 (might still be 0 if conformers happen to be similar) + assert rmsd >= 0 + + def test_remove_hs_option(self): + """Test that remove_hs option works.""" + mol = Chem.MolFromSmiles("C") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + + mol_copy = Chem.Mol(mol) + + # Both should return 0 for identical molecules + rmsd_with_hs_removed = get_rmsd(mol, mol_copy, remove_hs=True) + rmsd_without_hs_removed = get_rmsd(mol, mol_copy, remove_hs=False) + + assert abs(rmsd_with_hs_removed) < 1e-5 + assert abs(rmsd_without_hs_removed) < 1e-5 + + def test_mismatched_molecules_returns_inf(self): + """Mismatched molecules are incomparable and return float('inf'). + + An incomparable pair is treated as "distinct" (inf), matching + filter_unique, so a downstream `rmsd < threshold` check keeps the + structure instead of dropping it as a false duplicate. + """ + mol1 = Chem.MolFromSmiles("CCO") + mol1 = Chem.AddHs(mol1) + AllChem.EmbedMolecule(mol1, randomSeed=42) + + mol2 = Chem.MolFromSmiles("CCCC") + mol2 = Chem.AddHs(mol2) + AllChem.EmbedMolecule(mol2, randomSeed=42) + + # This should raise RuntimeError internally and return inf. + rmsd = get_rmsd(mol1, mol2) + assert rmsd == float("inf") + + def test_runtime_error_returns_inf(self, monkeypatch): + """A RuntimeError from the RMSD computation returns float('inf').""" + from Auto3D.utils import geometry as chem + + mol = Chem.MolFromSmiles("CCO") + mol = Chem.AddHs(mol) + AllChem.EmbedMolecule(mol, randomSeed=42) + mol_copy = Chem.Mol(mol) + + def _boom(*args, **kwargs): + raise RuntimeError("forced failure") + + monkeypatch.setattr(chem.rdMolAlign, "GetBestRMS", _boom) + assert chem.get_rmsd(mol, mol_copy) == float("inf") diff --git a/tests/test_chemistry.py b/tests/test_utils_molprops.py similarity index 65% rename from tests/test_chemistry.py rename to tests/test_utils_molprops.py index 3f0a6307..3d9ee2e3 100644 --- a/tests/test_chemistry.py +++ b/tests/test_utils_molprops.py @@ -1,9 +1,46 @@ #!/usr/bin/env python -"""Tests for Auto3D.utils.chemistry module.""" -import pytest +"""Tests for Auto3D.utils.molprops module.""" +from __future__ import annotations + from rdkit import Chem from Auto3D.constants import MAX_CONFORMERS_CAP +from Auto3D.utils.molprops import calculate_conformer_count, get_mol_charge + + +class TestGetMolCharge: + """Test the get_mol_charge function.""" + + def test_neutral_molecule(self): + """Test charge of a neutral molecule.""" + mol = Chem.MolFromSmiles("CCO") + assert get_mol_charge(mol) == 0 + + def test_cation(self): + """Test charge of a cation.""" + mol = Chem.MolFromSmiles("[NH4+]") + assert get_mol_charge(mol) == 1 + + def test_anion(self): + """Test charge of an anion.""" + mol = Chem.MolFromSmiles("[O-]") + assert get_mol_charge(mol) == -1 + + def test_doubly_charged_cation(self): + """Test charge of a doubly charged cation.""" + mol = Chem.MolFromSmiles("[Ca+2]") + assert get_mol_charge(mol) == 2 + + def test_zwitterion(self): + """Test charge of a zwitterion (net neutral).""" + # Glycine zwitterion + mol = Chem.MolFromSmiles("[NH3+]CC([O-])=O") + assert get_mol_charge(mol) == 0 + + def test_multiple_charges(self): + """Test molecule with multiple charged atoms.""" + mol = Chem.MolFromSmiles("[O-]C([O-])=O") # Carbonate + assert get_mol_charge(mol) == -2 class TestCalculateConformerCount: @@ -11,7 +48,6 @@ class TestCalculateConformerCount: def test_calculate_conformer_count_small_molecule(self): """Small molecule with few rotatable bonds should get reasonable count.""" - from Auto3D.utils.chemistry import calculate_conformer_count mol = Chem.MolFromSmiles("CCO") # ethanol - 0 rotatable bonds count = calculate_conformer_count(mol) @@ -20,7 +56,6 @@ def test_calculate_conformer_count_small_molecule(self): def test_calculate_conformer_count_flexible_molecule(self): """Flexible molecule should get higher conformer count.""" - from Auto3D.utils.chemistry import calculate_conformer_count mol = Chem.MolFromSmiles("CCCCCCCC") # octane - many rotatable bonds count = calculate_conformer_count(mol) @@ -28,7 +63,6 @@ def test_calculate_conformer_count_flexible_molecule(self): def test_calculate_conformer_count_respects_cap(self): """Very flexible molecules should be capped at MAX_CONFORMERS_CAP.""" - from Auto3D.utils.chemistry import calculate_conformer_count mol = Chem.MolFromSmiles("C" * 30) # very long chain count = calculate_conformer_count(mol) @@ -36,7 +70,6 @@ def test_calculate_conformer_count_respects_cap(self): def test_calculate_conformer_count_minimum_is_heavy_atoms(self): """Conformer count should be at least the number of heavy atoms.""" - from Auto3D.utils.chemistry import calculate_conformer_count mol = Chem.MolFromSmiles("C") # methane - single heavy atom count = calculate_conformer_count(mol) @@ -44,7 +77,6 @@ def test_calculate_conformer_count_minimum_is_heavy_atoms(self): def test_calculate_conformer_count_returns_int(self): """Result should always be an integer.""" - from Auto3D.utils.chemistry import calculate_conformer_count mol = Chem.MolFromSmiles("CCCCC") # pentane count = calculate_conformer_count(mol) @@ -52,7 +84,6 @@ def test_calculate_conformer_count_returns_int(self): def test_calculate_conformer_count_zero_rotatable_bonds(self): """Molecule with zero rotatable bonds should return at least heavy atom count.""" - from Auto3D.utils.chemistry import calculate_conformer_count mol = Chem.MolFromSmiles("c1ccccc1") # benzene - 0 rotatable bonds count = calculate_conformer_count(mol) @@ -61,7 +92,6 @@ def test_calculate_conformer_count_zero_rotatable_bonds(self): def test_calculate_conformer_count_molecule_with_hydrogens(self): """Function should work correctly with molecules that have explicit hydrogens.""" - from Auto3D.utils.chemistry import calculate_conformer_count mol = Chem.MolFromSmiles("CCO") mol_h = Chem.AddHs(mol) diff --git a/tests/test_utils_reconciliation.py b/tests/test_utils_reconciliation.py new file mode 100644 index 00000000..9a4fb42a --- /dev/null +++ b/tests/test_utils_reconciliation.py @@ -0,0 +1,255 @@ +"""Tests for Auto3D.utils.reconciliation module.""" +from pathlib import Path + +import pytest # noqa: F401 (used by the __main__ guard below) +from rdkit import Chem + +from Auto3D.utils.reconciliation import find_ids_not_in_sdf, find_smiles_not_in_sdf + +# Get the test files directory +TEST_DIR = Path(__file__).parent +FILES_DIR = TEST_DIR / "files" + + + +def _make_mol(name): + """Build a tiny named RDKit mol for SDF round-trips.""" + from rdkit import Chem + + mol = Chem.MolFromSmiles("C") + mol.SetProp("_Name", name) + return mol + + + +class TestNoneAndMalformedInputHardening: + """``None`` SDF records and lenient .smi lines must not crash reconciliation.""" + + def test_find_smiles_not_in_sdf_skips_none_records(self, tmp_path, monkeypatch): + """find_smiles_not_in_sdf must skip None SDF records.""" + from rdkit import Chem + + import Auto3D.utils.reconciliation as reconciliation + + valid = Chem.MolFromSmiles("C") + valid.SetProp("_Name", "mol_a") + monkeypatch.setattr( + reconciliation.Chem, "SDMolSupplier", lambda *a, **k: [valid, None] + ) + + smi = tmp_path / "in.smi" + smi.write_text("C mol_a\nCC mol_b\n") + sdf = tmp_path / "out.sdf" + sdf.write_text("placeholder") + + bad = find_smiles_not_in_sdf(str(smi), str(sdf)) + # mol_a is present (valid mol), mol_b is missing -> reported. + assert ("mol_b", "CC") in bad + assert all(mol_id != "mol_a" for mol_id, _ in bad) + def test_find_smiles_not_in_sdf_tolerates_blank_and_3token_lines( + self, tmp_path, monkeypatch + ): + """find_smiles_not_in_sdf tolerates blank and 3-token .smi lines.""" + from rdkit import Chem + + import Auto3D.utils.reconciliation as reconciliation + + valid = Chem.MolFromSmiles("C") + valid.SetProp("_Name", "mol_a") + monkeypatch.setattr( + reconciliation.Chem, "SDMolSupplier", lambda *a, **k: [valid] + ) + + smi = tmp_path / "in.smi" + # blank line + a 3-token line (first two tokens taken). + smi.write_text("C mol_a\n\nCC mol_b extra\n") + sdf = tmp_path / "out.sdf" + sdf.write_text("placeholder") + + bad = find_smiles_not_in_sdf(str(smi), str(sdf)) + assert ("mol_b", "CC") in bad + + +class TestFindSmilesNotInSdfTautStripping: + """C7: a decoded '@tautN' suffix must not cause a false "missing" report.""" + + def test_taut_suffixed_output_name_matches_base_smi_id(self, tmp_path): + """decode_ids keeps 'id@tautN' on tautomer conformers; the .smi only + has the base id, so find_smiles_not_in_sdf must strip the suffix + before comparing or every tautomer-derived molecule is misreported.""" + smi = tmp_path / "in.smi" + smi.write_text("CCO mol_a\n") + + sdf = tmp_path / "out.sdf" + writer = Chem.SDWriter(str(sdf)) + writer.write(_make_mol("mol_a@taut0")) + writer.close() + + bad = find_smiles_not_in_sdf(str(smi), str(sdf)) + assert bad == [], f"mol_a wrongly reported missing: {bad}" + + +class TestFindIdsNotInSdf: + """find_ids_not_in_sdf: the SDF-input counterpart to find_smiles_not_in_sdf.""" + + def test_missing_id_is_reported(self, tmp_path): + """An id present in the source SDF but absent from the output SDF is reported.""" + source = tmp_path / "source.sdf" + writer = Chem.SDWriter(str(source)) + for name in ["mol_a", "mol_b"]: + writer.write(_make_mol(name)) + writer.close() + + out = tmp_path / "out.sdf" + writer = Chem.SDWriter(str(out)) + writer.write(_make_mol("mol_a")) # mol_b never produced a structure + writer.close() + + bad = find_ids_not_in_sdf(str(source), str(out)) + assert bad == ["mol_b"] + + def test_no_missing_ids_returns_empty_list(self, tmp_path): + """Every source id present in the output -> nothing reported.""" + source = tmp_path / "source.sdf" + writer = Chem.SDWriter(str(source)) + for name in ["mol_a", "mol_b"]: + writer.write(_make_mol(name)) + writer.close() + + out = tmp_path / "out.sdf" + writer = Chem.SDWriter(str(out)) + for name in ["mol_b", "mol_a"]: + writer.write(_make_mol(name)) + writer.close() + + assert find_ids_not_in_sdf(str(source), str(out)) == [] + + def test_taut_suffixed_output_name_matches_base_id(self, tmp_path): + """Same '@tautN' stripping as find_smiles_not_in_sdf, for SDF input.""" + source = tmp_path / "source.sdf" + writer = Chem.SDWriter(str(source)) + writer.write(_make_mol("mol_a")) + writer.close() + + out = tmp_path / "out.sdf" + writer = Chem.SDWriter(str(out)) + writer.write(_make_mol("mol_a@taut1")) + writer.close() + + assert find_ids_not_in_sdf(str(source), str(out)) == [] + + def test_an_unreadable_record_is_reported_on_the_input_side_only( + self, tmp_path, monkeypatch + ): + """The two sides of the comparison are not symmetric, and must not be. + + This test used to assert ``== []`` for a source file containing an + unreadable record, under the heading "must not crash or miscount". The + empty list *was* the miscount: an input molecule the pipeline never saw + was reported by nothing, and the run exited 0 claiming completeness. + + The asymmetry is the point: + + * **input side** -- a record that cannot be read is a molecule the user + supplied and did not get back. It is reported, by position, since it has + no ``_Name`` to report by. + * **output side** -- a record that cannot be read yields no name to match + against, and there is nothing better to do than skip it. Reporting it + would invent a *missing input* out of an unreadable output. + """ + import Auto3D.utils.reconciliation as reconciliation + + calls = {"n": 0} + + def fake_supplier(*a, **k): + calls["n"] += 1 + if calls["n"] == 1: + return [_make_mol("mol_a"), None] # source: one good, one unreadable + return [None, _make_mol("mol_a")] # output: mol_a made it + + monkeypatch.setattr(reconciliation.Chem, "SDMolSupplier", fake_supplier) + + source = tmp_path / "source.sdf" + source.write_text("placeholder") + out = tmp_path / "out.sdf" + out.write_text("placeholder") + + assert find_ids_not_in_sdf(str(source), str(out)) == [ + reconciliation.UNPARSEABLE_RECORD_ID.format(index=1) + ] + + +class TestReconciliationSeesUnparseableInputRecords: + """A record Auto3D could not read must not vanish from the accounting. + + ``encode_ids`` skips an unparseable SDF record with a warning, so it never + enters the run. ``find_ids_not_in_sdf`` then built its expected-ID list by + reading **the same source SDF** and skipping the same record -- so it was in + neither ``source_ids`` nor the output, could not appear in ``failures``, and + ``_exit_if_incomplete`` saw ``failed_count == 0``. The run printed a success + summary and exited **0** having processed fewer molecules than the file + contained, which is precisely what the C7 reconciliation exists to prevent. + + Only the SDF path is affected. ``encode_ids`` reads ``.smi`` input with + ``on_malformed="raise"``, so a malformed SMILES line aborts the run with + ``InputValidationError`` long before reconciliation -- the same blindness + cannot be reached through that door. + """ + + @staticmethod + def _unparseable_record(name: str) -> str: + """A molblock RDKit rejects: the counts line is corrupted. + + Built from a real molblock so only the one line under test is invalid; + a wholly invented block could fail for an unrelated reason. + """ + mol = _make_mol(name) + lines = Chem.MolToMolBlock(mol).splitlines() + lines[3] = "!! corrupted counts line !!" + return "\n".join(lines) + + def test_an_unparseable_source_record_is_reported_as_a_failure(self, tmp_path): + source = tmp_path / "source.sdf" + good = Chem.MolToMolBlock(_make_mol("mol_a")) + source.write_text( + good + "$$$$\n" + self._unparseable_record("mol_b") + "\n$$$$\n" + ) + # Confirm the premise: RDKit reads one molecule and one None. + parsed = list(Chem.SDMolSupplier(str(source), removeHs=False)) + assert [m is None for m in parsed] == [False, True], "test premise" + + out = tmp_path / "out.sdf" + writer = Chem.SDWriter(str(out)) + writer.write(_make_mol("mol_a")) # the parseable one succeeded + writer.close() + + missing = find_ids_not_in_sdf(str(source), str(out)) + + assert len(missing) == 1, ( + f"a source record the pipeline could not read was left out of the " + f"accounting entirely, so the run would exit 0 claiming every input " + f"was processed; got {missing}" + ) + assert "1" in missing[0], ( + f"the report must say which record could not be read, got {missing[0]!r}" + ) + + def test_a_clean_source_file_still_reports_nothing(self, tmp_path): + """The new branch must not manufacture failures for a healthy file.""" + source = tmp_path / "source.sdf" + writer = Chem.SDWriter(str(source)) + for name in ("mol_a", "mol_b"): + writer.write(_make_mol(name)) + writer.close() + + out = tmp_path / "out.sdf" + writer = Chem.SDWriter(str(out)) + for name in ("mol_a", "mol_b"): + writer.write(_make_mol(name)) + writer.close() + + assert find_ids_not_in_sdf(str(source), str(out)) == [] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_utils_sdf_io.py b/tests/test_utils_sdf_io.py new file mode 100644 index 00000000..67f9bdae --- /dev/null +++ b/tests/test_utils_sdf_io.py @@ -0,0 +1,321 @@ +"""Tests for Auto3D.utils.sdf_io module.""" +from pathlib import Path + +import pytest # noqa: F401 (used by the __main__ guard below) +from rdkit import Chem + +from Auto3D.utils.sdf_io import ( + SDF2chunks, + count_sdf, + guess_file_type, + reorder_sdf, +) + +# Get the test files directory +TEST_DIR = Path(__file__).parent +FILES_DIR = TEST_DIR / "files" + + + +def _make_mol(name): + """Build a tiny named RDKit mol for SDF round-trips.""" + + mol = Chem.MolFromSmiles("C") + mol.SetProp("_Name", name) + return mol + + +class TestGuessFileType: + """Tests for guess_file_type function.""" + + def test_smi_extension(self): + """Test detection of .smi files.""" + assert guess_file_type("molecules.smi") == "smi" + assert guess_file_type("/path/to/input.smi") == "smi" + + def test_sdf_extension(self): + """Test detection of .sdf files.""" + assert guess_file_type("molecules.sdf") == "sdf" + assert guess_file_type("/data/output/result.sdf") == "sdf" + + def test_mol2_extension(self): + """Test detection of .mol2 files.""" + assert guess_file_type("molecule.mol2") == "mol2" + + def test_xyz_extension(self): + """Test detection of .xyz files.""" + assert guess_file_type("geometry.xyz") == "xyz" + + def test_complex_path(self): + """Test with complex file paths.""" + assert guess_file_type("/home/user/data.2024/molecules.sdf") == "sdf" + assert guess_file_type("./relative/path/file.smi") == "smi" + + def test_no_extension(self): + """Test file without extension returns empty string.""" + assert guess_file_type("filename") == "" + + def test_hidden_file(self): + """Test hidden files with extension.""" + assert guess_file_type(".hidden.sdf") == "sdf" + + +class TestSDF2chunks: + """Tests for SDF2chunks function.""" + + def test_splits_sdf_into_chunks(self): + """Test that SDF file is split into molecule chunks.""" + sdf_path = str(FILES_DIR / "example.sdf") + + chunks = SDF2chunks(sdf_path) + + # example.sdf has 2 molecules + assert len(chunks) == 2 + + # Each chunk should end with $$$$ + for chunk in chunks: + assert chunk[-1].strip() == "$$$$" + + def test_chunk_contains_molecule_lines(self): + """Test that chunks contain all molecule lines.""" + sdf_path = str(FILES_DIR / "example.sdf") + + chunks = SDF2chunks(sdf_path) + + # First chunk should start with molecule name + assert chunks[0][0].strip() == "mol1" + assert chunks[1][0].strip() == "mol2" + + def test_preserves_all_content(self): + """Test that all content from original file is preserved.""" + sdf_path = str(FILES_DIR / "example.sdf") + + chunks = SDF2chunks(sdf_path) + + # Reconstruct file from chunks + reconstructed = "".join(line for chunk in chunks for line in chunk) + + with open(sdf_path) as f: + original = f.read() + + assert reconstructed == original + + +class TestReorderSdf: + """Tests for reorder_sdf function.""" + + def test_reorder_sdf_from_smi(self, tmp_path): + """Test reordering SDF file based on SMILES file order.""" + + # Create source SMILES file with specific order + smi_file = tmp_path / "source.smi" + smi_file.write_text("CCO mol_b\nCC mol_a\nCCC mol_c\n") + + # Create SDF file with different order + sdf_file = tmp_path / "mols.sdf" + writer = Chem.SDWriter(str(sdf_file)) + for name in ["mol_a", "mol_c", "mol_b"]: + mol = Chem.MolFromSmiles("C") + mol.SetProp("_Name", name) + writer.write(mol) + writer.close() + + # Reorder + result = reorder_sdf(str(sdf_file), str(smi_file)) + + # Verify order matches source + assert len(result) == 3 + assert result[0].GetProp("_Name") == "mol_b" + assert result[1].GetProp("_Name") == "mol_a" + assert result[2].GetProp("_Name") == "mol_c" + + def test_reorder_sdf_from_sdf(self, tmp_path): + """Test reordering SDF file based on another SDF file order.""" + + # Create source SDF file with specific order + source_sdf = tmp_path / "source.sdf" + writer = Chem.SDWriter(str(source_sdf)) + for name in ["mol_x", "mol_y", "mol_z"]: + mol = Chem.MolFromSmiles("C") + mol.SetProp("_Name", name) + writer.write(mol) + writer.close() + + # Create target SDF file with different order + target_sdf = tmp_path / "target.sdf" + writer = Chem.SDWriter(str(target_sdf)) + for name in ["mol_z", "mol_x", "mol_y"]: + mol = Chem.MolFromSmiles("C") + mol.SetProp("_Name", name) + writer.write(mol) + writer.close() + + # Reorder + result = reorder_sdf(str(target_sdf), str(source_sdf)) + + # Verify order matches source + assert len(result) == 3 + assert result[0].GetProp("_Name") == "mol_x" + assert result[1].GetProp("_Name") == "mol_y" + assert result[2].GetProp("_Name") == "mol_z" + + def test_reorder_sdf_with_tautomers(self, tmp_path): + """Test reordering handles tautomer IDs correctly.""" + + # Create source SMILES file + smi_file = tmp_path / "source.smi" + smi_file.write_text("CCO mol1\nCC mol2\n") + + # Create SDF file with tautomer variants + sdf_file = tmp_path / "mols.sdf" + writer = Chem.SDWriter(str(sdf_file)) + for name in ["mol2@taut1", "mol1@taut1", "mol1@taut2"]: + mol = Chem.MolFromSmiles("C") + mol.SetProp("_Name", name) + writer.write(mol) + writer.close() + + # Reorder + result = reorder_sdf(str(sdf_file), str(smi_file)) + + # Verify mol1 variants come before mol2 variants + assert len(result) == 3 + # mol1 should be first (2 tautomers) + assert "mol1" in result[0].GetProp("_Name") + assert "mol1" in result[1].GetProp("_Name") + # mol2 should be last + assert "mol2" in result[2].GetProp("_Name") + + def test_reorder_sdf_unsupported_format(self, tmp_path, caplog): + """Test that unsupported format returns None.""" + import logging + + xyz_file = tmp_path / "source.xyz" + xyz_file.write_text("invalid") + + sdf_file = tmp_path / "mols.sdf" + sdf_file.write_text("dummy") + + with caplog.at_level(logging.WARNING): + result = reorder_sdf(str(sdf_file), str(xyz_file)) + + assert result is None + assert "Unsupported file format" in caplog.text + + +class TestNoneMolHardening: + """FIX 1: None records yielded by SDMolSupplier must not crash these helpers. + + A single unparseable SDF record makes SDMolSupplier yield ``None``. The + iterating helpers previously called ``mol.GetProp(...)`` / ``mol.GetNumAtoms()`` + on it and raised ``AttributeError``. They must skip ``None`` instead. + """ + + def test_count_sdf_skips_none_records(self, tmp_path, monkeypatch): + """count_sdf must not count (or crash on) a None record.""" + + import Auto3D.utils.sdf_io as sdf_io + + valid = _make_mol("mol_a") + monkeypatch.setattr( + sdf_io.Chem, "SDMolSupplier", lambda *a, **k: [valid, None] + ) + + sdf = tmp_path / "mols.sdf" + sdf.write_text("placeholder") # path only needs to exist for the call + + assert count_sdf(str(sdf)) == 1 + def test_reorder_sdf_skips_none_records(self, tmp_path, monkeypatch): + """reorder_sdf must skip None records in the target SDF.""" + + import Auto3D.utils.sdf_io as sdf_io + + smi = tmp_path / "source.smi" + smi.write_text("C mol_a\nC mol_b\n") + + valid_a = _make_mol("mol_a") + valid_b = _make_mol("mol_b") + monkeypatch.setattr( + sdf_io.Chem, "SDMolSupplier", lambda *a, **k: [valid_a, None, valid_b] + ) + + sdf = tmp_path / "target.sdf" + sdf.write_text("placeholder") + + result = reorder_sdf(str(sdf), str(smi)) + names = [m.GetProp("_Name") for m in result] + assert names == ["mol_a", "mol_b"] + + +class TestReorderSdfDataPreservation: + """FIX 2: reorder_sdf must not drop unmatched molecules or truncate input.""" + + def test_unmatched_mol_is_preserved(self, tmp_path): + """A mol whose id is not in the source must still survive to disk.""" + + smi = tmp_path / "source.smi" + # source lists only mol_a and mol_b; mol_c is unmatched. + smi.write_text("C mol_a\nC mol_b\n") + + sdf = tmp_path / "target.sdf" + writer = Chem.SDWriter(str(sdf)) + for name in ["mol_c", "mol_b", "mol_a"]: + writer.write(_make_mol(name)) + writer.close() + + result = reorder_sdf(str(sdf), str(smi)) + + # All three mols preserved (no silent data loss). + result_names = [m.GetProp("_Name") for m in result] + assert set(result_names) == {"mol_a", "mol_b", "mol_c"} + assert len(result_names) == 3 + # Matched ids appear first, in source order. + assert result_names[0] == "mol_a" + assert result_names[1] == "mol_b" + + # And the on-disk file must contain all three as well. + on_disk = [m.GetProp("_Name") for m in Chem.SDMolSupplier(str(sdf))] + assert set(on_disk) == {"mol_a", "mol_b", "mol_c"} + + def test_normal_all_matched_ordering_unchanged(self, tmp_path): + """When every id is matched, ordering is exactly the source order.""" + + smi = tmp_path / "source.smi" + smi.write_text("C mol_b\nC mol_a\nC mol_c\n") + + sdf = tmp_path / "target.sdf" + writer = Chem.SDWriter(str(sdf)) + for name in ["mol_a", "mol_c", "mol_b"]: + writer.write(_make_mol(name)) + writer.close() + + result = reorder_sdf(str(sdf), str(smi)) + names = [m.GetProp("_Name") for m in result] + assert names == ["mol_b", "mol_a", "mol_c"] + + +class TestSDF2chunksTrailingRecord: + """FIX 3: a final record lacking the $$$$ terminator must not be dropped.""" + + def test_trailing_record_without_terminator_preserved(self, tmp_path): + """SDF2chunks keeps a terminator-less trailing record as its own chunk.""" + sdf = tmp_path / "ragged.sdf" + # First record has $$$$; second record lacks it. + sdf.write_text( + "mol1\n line1\n$$$$\n" + "mol2\n line2\n line3\n" + ) + + chunks = SDF2chunks(str(sdf)) + + assert len(chunks) == 2 + assert chunks[0][0].strip() == "mol1" + # The trailing record's lines must be present in the final chunk. + assert chunks[1][0].strip() == "mol2" + joined = "".join(chunks[1]) + assert "line2" in joined + assert "line3" in joined + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_utils_smi_io.py b/tests/test_utils_smi_io.py new file mode 100644 index 00000000..5f487a9f --- /dev/null +++ b/tests/test_utils_smi_io.py @@ -0,0 +1,363 @@ +"""Tests for Auto3D.utils.smi_io module.""" +from pathlib import Path + +import pytest +from rdkit import Chem # noqa: F401 (kept for parity with the sibling io tests) + +from Auto3D.utils.smi_io import ( + combine_smi, + hash_enumerated_smi_IDs, + hash_taut_smi, + iter_smi_records, + smiles2smi, +) + +# Get the test files directory +TEST_DIR = Path(__file__).parent +FILES_DIR = TEST_DIR / "files" + + +class TestSmiles2Smi: + """Tests for smiles2smi function.""" + + def test_creates_file_with_inchikeys(self, tmp_path): + """smiles2smi should create a .smi file with SMILES and InChIKey IDs.""" + smiles = ["CCO", "CCC"] + output = tmp_path / "test.smi" + + result = smiles2smi(smiles, str(output)) + + assert result == str(output) + assert output.exists() + content = output.read_text() + lines = content.strip().split('\n') + assert len(lines) == 2 + # Each line should have SMILES and InChIKey + for line in lines: + parts = line.split() + assert len(parts) == 2 + + def test_returns_output_path(self, tmp_path): + """smiles2smi should return the output file path.""" + smiles = ["CCO"] + output = tmp_path / "output.smi" + + result = smiles2smi(smiles, str(output)) + + assert result == str(output) + + def test_inchikey_format(self, tmp_path): + """InChIKeys should have the standard 27-character format.""" + smiles = ["CCO"] + output = tmp_path / "test.smi" + + smiles2smi(smiles, str(output)) + + content = output.read_text().strip() + parts = content.split() + inchikey = parts[1] + # InChIKey format: 14 chars + hyphen + 10 chars + hyphen + 1 char = 27 chars + assert len(inchikey) == 27 + assert inchikey.count('-') == 2 + + def test_preserves_smiles_string(self, tmp_path): + """Original SMILES strings should be preserved in output.""" + smiles = ["C#N", "C=C", "[NH4+]"] + output = tmp_path / "test.smi" + + smiles2smi(smiles, str(output)) + + content = output.read_text() + for smi in smiles: + assert smi in content + + def test_empty_list(self, tmp_path): + """Empty input list should create empty file.""" + output = tmp_path / "test.smi" + + result = smiles2smi([], str(output)) + + assert result == str(output) + assert output.exists() + assert output.read_text() == "" + + def test_colliding_inchikeys_get_distinct_ids(self, tmp_path): + """Two inputs that share an InChIKey must keep distinct IDs. + + The same molecule written two ways (here benzene) yields one InChIKey; + without disambiguation reorder_sdf would collapse the duplicate IDs and + silently drop the second input. Each input must get its own line/ID. + """ + output = tmp_path / "test.smi" + + smiles2smi(["c1ccccc1", "C1=CC=CC=C1"], str(output)) + + lines = output.read_text().strip().split("\n") + assert len(lines) == 2 + ids = [line.split()[1] for line in lines] + assert ids[0] != ids[1], "colliding InChIKeys must be disambiguated" + # First keeps the bare InChIKey; the repeat is suffixed. + assert ids[1] == f"{ids[0]}_2" + + def test_distinct_inputs_keep_bare_inchikeys(self, tmp_path): + """Non-colliding inputs must keep their plain InChIKey IDs (no suffix).""" + output = tmp_path / "test.smi" + + smiles2smi(["CCO", "CCC"], str(output)) + + ids = [line.split()[1] for line in output.read_text().strip().split("\n")] + assert ids[0] != ids[1] + assert all("_" not in i for i in ids) + + +class TestCombineSmiOrderPreservingDedup: + """Tests for combine_smi (order-preserving dedup). + + Named distinctly from the ``TestCombineSmi`` class below -- both defined + ``TestCombineSmi`` until a lint audit found the second definition was + silently shadowing this one, so pytest never collected the test below. + """ + + def test_preserves_order_and_dedups(self, tmp_path): + f1 = tmp_path / "a.smi" + f2 = tmp_path / "b.smi" + f1.write_text("CCO ethanol\nCCC propane\n") + f2.write_text("CCC propane\nCCCC butane\n") # propane duplicated + out = tmp_path / "combined.smi" + + combine_smi([str(f1), str(f2)], str(out)) + + lines = out.read_text().strip().split("\n") + # Deduped (propane once) and in first-seen input order. + assert lines == ["CCO ethanol", "CCC propane", "CCCC butane"] + + +class TestHashEnumeratedSmiIDs: + """Tests for hash_enumerated_smi_IDs function.""" + + def test_basic_hashing(self, tmp_path): + """Test basic hashing with simple SMILES file.""" + input_file = tmp_path / "input.smi" + output_file = tmp_path / "output.smi" + + # Create input file with unsorted IDs + input_file.write_text("CCO mol_b\nCC mol_a\nCCC mol_c\n") + + hash_enumerated_smi_IDs(str(input_file), str(output_file)) + + # Read and verify output + lines = output_file.read_text().strip().split("\n") + assert len(lines) == 3 + # Should be sorted by ID + assert "mol_a" in lines[0] + assert "mol_b" in lines[1] + assert "mol_c" in lines[2] + + def test_duplicate_id_handling(self, tmp_path): + """Test that duplicate IDs get '_0' suffix.""" + input_file = tmp_path / "input.smi" + output_file = tmp_path / "output.smi" + + # Create input file with duplicate IDs + input_file.write_text("CCO mol1\nCC mol1\nCCC mol1\n") + + hash_enumerated_smi_IDs(str(input_file), str(output_file)) + + lines = output_file.read_text().strip().split("\n") + assert len(lines) == 3 + + # Check that duplicates were renamed + ids = [line.split()[1] for line in lines] + assert "mol1" in ids + assert "mol1_0" in ids + assert "mol1_0_0" in ids + + def test_preserves_smiles(self, tmp_path): + """Test that SMILES strings are preserved correctly.""" + input_file = tmp_path / "input.smi" + output_file = tmp_path / "output.smi" + + input_file.write_text("C#N id1\nC=C id2\n") + + hash_enumerated_smi_IDs(str(input_file), str(output_file)) + + content = output_file.read_text() + assert "C#N" in content + assert "C=C" in content + + +class TestHashTautSmi: + """Tests for hash_taut_smi function.""" + + def test_tautomer_suffix_added(self, tmp_path): + """Test that @taut suffix is added to IDs.""" + input_file = tmp_path / "input.smi" + output_file = tmp_path / "output.smi" + + input_file.write_text("CCO mol1\nCC mol2\n") + + hash_taut_smi(str(input_file), str(output_file)) + + content = output_file.read_text() + assert "@taut" in content + + def test_incremental_taut_suffix(self, tmp_path): + """Test that duplicate base IDs get incrementing taut numbers.""" + input_file = tmp_path / "input.smi" + output_file = tmp_path / "output.smi" + + # Same ID for multiple SMILES + input_file.write_text("CCO mol1\nCC mol1\n") + + hash_taut_smi(str(input_file), str(output_file)) + + lines = output_file.read_text().strip().split("\n") + ids = [line.split()[1] for line in lines] + + # Should have different taut numbers + assert len(set(ids)) == 2 + assert all("@taut" in id for id in ids) + + +class TestCombineSmi: + """Tests for combine_smi function.""" + + def test_combines_files(self, tmp_path): + """Test that multiple SMILES files are combined.""" + file1 = tmp_path / "file1.smi" + file2 = tmp_path / "file2.smi" + output = tmp_path / "combined.smi" + + file1.write_text("CCO mol1\nCC mol2\n") + file2.write_text("CCC mol3\nCCCC mol4\n") + + combine_smi([str(file1), str(file2)], str(output)) + + content = output.read_text() + assert "mol1" in content + assert "mol2" in content + assert "mol3" in content + assert "mol4" in content + + def test_removes_duplicates(self, tmp_path): + """Test that duplicate entries are removed.""" + file1 = tmp_path / "file1.smi" + file2 = tmp_path / "file2.smi" + output = tmp_path / "combined.smi" + + file1.write_text("CCO mol1\n") + file2.write_text("CCO mol1\n") # Same entry + + combine_smi([str(file1), str(file2)], str(output)) + + lines = output.read_text().strip().split("\n") + assert len(lines) == 1 + + def test_ignores_blank_lines(self, tmp_path): + """Test that blank lines are ignored.""" + file1 = tmp_path / "file1.smi" + output = tmp_path / "combined.smi" + + file1.write_text("CCO mol1\n\n\nCC mol2\n \n") + + combine_smi([str(file1)], str(output)) + + lines = output.read_text().strip().split("\n") + assert len(lines) == 2 + + +class TestSmiles2SmiInvalidInput: + """FIX 4: smiles2smi must raise a clear error on an invalid SMILES.""" + + def test_invalid_smiles_raises_input_validation_error(self, tmp_path): + """An unparseable SMILES raises InputValidationError naming the SMILES.""" + from Auto3D.exceptions import InputValidationError + + out = tmp_path / "out.smi" + with pytest.raises(InputValidationError, match=r"C\(C"): + smiles2smi(["CCO", "C(C"], str(out)) + + +class TestHashHelpersBlankLines: + """FIX 5: blank / malformed lines must not crash the hashing helpers.""" + + def test_hash_enumerated_skips_blank_and_extra_token_lines(self, tmp_path): + """hash_enumerated_smi_IDs tolerates blank lines and extra tokens.""" + inp = tmp_path / "in.smi" + inp.write_text("CCO mol1\n\n \nCC mol2 extra_token\n") + out = tmp_path / "out.smi" + + # Must not raise ValueError. + hash_enumerated_smi_IDs(str(inp), str(out)) + + lines = [ln for ln in out.read_text().splitlines() if ln.strip()] + ids = [ln.split()[1] for ln in lines] + assert "mol1" in ids + assert "mol2" in ids + + def test_hash_taut_skips_blank_and_extra_token_lines(self, tmp_path): + """hash_taut_smi tolerates blank lines and extra tokens.""" + inp = tmp_path / "in.smi" + inp.write_text("CCO mol1\n\nCC mol2 extra_token\n") + out = tmp_path / "out.smi" + + hash_taut_smi(str(inp), str(out)) + + lines = [ln for ln in out.read_text().splitlines() if ln.strip()] + assert len(lines) == 2 + assert all("@taut" in ln.split()[1] for ln in lines) + +class TestIterSmiRecords: + """FIX A: shared lenient .smi parser used by all 7 call sites.""" + + def test_blank_lines_skipped(self, tmp_path): + """Blank and whitespace-only lines yield no records.""" + p = tmp_path / "in.smi" + p.write_text("CCO mol1\n\n \nCC mol2\n") + records = list(iter_smi_records(str(p))) + assert [(s, i) for _ln, s, i in records] == [("CCO", "mol1"), ("CC", "mol2")] + # line_no is 1-based and reflects the original line position. + assert records[0][0] == 1 + assert records[1][0] == 4 + + def test_three_token_line_yields_first_two(self, tmp_path): + """A 3-token line yields only the first two tokens (extras ignored).""" + p = tmp_path / "in.smi" + p.write_text("CCN extra_a extra_b\n") + records = list(iter_smi_records(str(p))) + assert len(records) == 1 + line_no, smiles, mol_id = records[0] + assert (smiles, mol_id) == ("CCN", "extra_a") + + def test_on_malformed_skip_skips_one_token_line_with_warning( + self, tmp_path, caplog + ): + """on_malformed='skip' (default) skips a 1-token line and warns.""" + import logging + + p = tmp_path / "in.smi" + p.write_text("CCO mol1\nC1CCCCC1\nCC mol2\n") + with caplog.at_level(logging.WARNING): + records = list(iter_smi_records(str(p), on_malformed="skip")) + assert [(s, i) for _ln, s, i in records] == [("CCO", "mol1"), ("CC", "mol2")] + assert any("failed to parse" in r.message for r in caplog.records) + + def test_on_malformed_raise_raises_on_one_token_line(self, tmp_path): + """on_malformed='raise' raises InputValidationError naming the line.""" + from Auto3D.exceptions import InputValidationError + + p = tmp_path / "in.smi" + p.write_text("CCO mol1\nC1CCCCC1\n") + with pytest.raises(InputValidationError, match="Line 2"): + list(iter_smi_records(str(p), on_malformed="raise")) + + def test_invalid_on_malformed_value_raises(self, tmp_path): + """An unknown on_malformed value raises ValueError.""" + p = tmp_path / "in.smi" + p.write_text("CCO mol1\n") + with pytest.raises(ValueError, match="on_malformed"): + list(iter_smi_records(str(p), on_malformed="bogus")) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_utils_validation.py b/tests/test_utils_validation.py index 35c7a7e2..a1643907 100644 --- a/tests/test_utils_validation.py +++ b/tests/test_utils_validation.py @@ -6,7 +6,8 @@ from rdkit import Chem from Auto3D.config import Auto3DOptions -from Auto3D.utils.chemistry import check_connectivity, filter_unique +from Auto3D.filtering import filter_unique +from Auto3D.utils.connectivity import check_connectivity from Auto3D.utils.validation import ( check_input, check_sdf_format, diff --git a/tests/test_validate_run_parity.py b/tests/test_validate_run_parity.py index efe79da1..ad8a3762 100644 --- a/tests/test_validate_run_parity.py +++ b/tests/test_validate_run_parity.py @@ -2,7 +2,7 @@ """M25 parity: `auto3d validate` must reject exactly what the runner rejects. Before this fix, cli.commands.validate.validate_smiles_file did not require an -ID column (it took parts[0] with no length check), while file_ops.encode_ids +ID column (it took parts[0] with no length check), while id_mapping.encode_ids (via iter_smi_records, on_malformed="raise") always has -- so a SMILES-only file passed `auto3d validate` and then failed the run, whose own error hint told the user to run the validator that had just approved it. The two also @@ -20,7 +20,7 @@ from Auto3D.cli.commands.validate import validate_smiles_file from Auto3D.exceptions import InputValidationError -from Auto3D.utils.file_ops import encode_ids +from Auto3D.id_mapping import encode_ids def _validator_accepts(path) -> bool: diff --git a/tests/test_workflow.py b/tests/test_workflow.py index 1bed7b59..3583f166 100644 --- a/tests/test_workflow.py +++ b/tests/test_workflow.py @@ -826,15 +826,17 @@ def test_sdf_input_reports_missing_id_and_sets_failures(self, tmp_path): orch._finalize_output(start_time=0.0) assert orch.failures == ["mol_c"], orch.failures - def test_workflow_uses_the_canonical_file_ops_reconciliation_functions(self): + def test_workflow_uses_the_canonical_reconciliation_functions(self): """Guard against a regression to a hand-rolled duplicate: workflow.py - must call the exact functions tested in test_utils_file_ops.py, not a + must call the exact functions tested in test_utils_reconciliation.py, not a reimplementation that could silently diverge from them.""" - import Auto3D.utils.file_ops as file_ops + import Auto3D.utils.reconciliation as reconciliation import Auto3D.workflow as workflow - assert workflow.find_smiles_not_in_sdf is file_ops.find_smiles_not_in_sdf - assert workflow.find_ids_not_in_sdf is file_ops.find_ids_not_in_sdf + assert ( + workflow.find_smiles_not_in_sdf is reconciliation.find_smiles_not_in_sdf + ) + assert workflow.find_ids_not_in_sdf is reconciliation.find_ids_not_in_sdf def test_main_propagates_orchestrator_failures_into_workflow_result(monkeypatch, tmp_path): @@ -1015,7 +1017,7 @@ def test_the_parallel_embed_path_names_a_species_it_produced_nothing_for( signal: a message logged inside a ProcessPoolExecutor worker depends on that child's logging configuration, and this one does not. """ - from Auto3D.isomers.parallel_embed import embed_conformers_parallel + from Auto3D.embedding import embed_conformers_parallel with caplog.at_level(logging.WARNING): results = list( @@ -1034,7 +1036,7 @@ def test_the_parallel_embed_path_names_a_species_it_produced_nothing_for( def test_a_species_that_embeds_normally_is_not_warned_about(self, caplog): """The new branch must not fire for a molecule that worked.""" - from Auto3D.isomers.parallel_embed import embed_conformers_parallel + from Auto3D.embedding import embed_conformers_parallel with caplog.at_level(logging.WARNING): results = list( From 2ba815287a7eadbc8d6fda4325ccbeb904ec9b32 Mon Sep 17 00:00:00 2001 From: isayev Date: Tue, 4 Aug 2026 10:29:04 -0400 Subject: [PATCH 2/2] test: make the malformed-record test possible to pass The assertion was strengthened from "any valid molecule produced a result" to "both did", which is the right assertion -- a batch that aborts right after the corrupt record satisfies the weaker one. But it was derived from reading the source, and the test needs a real NNP, so it could not be run where it was written. CI's slow tier then failed with {'ethanol'} != {'ethanol', 'propanol'}. The code was never the problem. The input was. RDKit's SDMolSupplier, given "this is not a molecule\n$$$$\n" between two valid records, logs "moving to the beginning of the next molecule" and consumes the following record while resynchronizing: the supplier yields [ethanol, None] and propanol is never handed over at all. No implementation of calc_thermo could have produced a propanol result from that file. The corrupt block is now a well-delimited record -- header, counts line promising two atoms, garbage where the coordinates belong, M END, $$$$ -- so the supplier yields exactly [ethanol, None, propanol]. Measured against this repo's RDKit 2025.09.6, along with a check that iter_thermo_records skips the None and keeps both valid records, so the guarded path is confirmed to fire without loading a model. The docstring now records both supplier behaviors the input depends on, since each one silently makes the test vacuous in a different way. --- tests/test_thermo_reference.py | 44 ++++++++++++++++++++++++---------- 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/tests/test_thermo_reference.py b/tests/test_thermo_reference.py index 89e4d9d2..411155c1 100644 --- a/tests/test_thermo_reference.py +++ b/tests/test_thermo_reference.py @@ -35,15 +35,23 @@ class TestBatchRobustness: def test_malformed_record_does_not_abort_the_batch(self, job_dir): """A None record between two valid ones must be skipped, not crash. - The corrupt block must sit BETWEEN two valid records, not after the - last one: verified against this repo's RDKit (2025.09.6), a forward - ``SDMolSupplier`` iterator (what ``list(...)`` and ``calc_thermo`` - itself use) only ever surfaces an explicit ``None`` for a corrupt - record when a further valid record follows it in the file. A corrupt - trailing record with nothing after it is silently dropped without - producing a ``None`` at all, so the code path under test (an - unguarded ``None.GetConformer()``) would never fire and the test - would pass today for the wrong reason. + Two properties of RDKit's ``SDMolSupplier`` shape this input, both + measured against this repo's version (2025.09.6) rather than assumed: + + 1. The corrupt block must sit BETWEEN two valid records, not after the + last one. A forward iterator (what ``list(...)`` and ``calc_thermo`` + both use) only surfaces an explicit ``None`` for a corrupt record + when a further valid record follows it; a corrupt trailing record is + dropped silently, so the guarded path would never fire. + + 2. The corrupt block must be a **well-delimited record** -- header, + counts line, garbage where coordinates belong, ``M END``, ``$$$$``. + Loose garbage does not work: on ``"this is not a molecule\\n$$$$\\n"`` + the supplier logs "moving to the beginning of the next molecule" and + **consumes the following record while resynchronizing**, yielding + ``[ethanol, None]`` with propanol never handed over at all. The + assertion below would then be impossible to satisfy for a reason that + has nothing to do with the code under test. """ from Auto3D.ASE.thermo import calc_thermo @@ -52,11 +60,21 @@ def test_malformed_record_does_not_abort_the_batch(self, job_dir): _write_mol(good1, smiles="CCO", name="ethanol") _write_mol(good2, smiles="CCCO", name="propanol") - # Sandwich a deliberately corrupt record between two valid ones. - combined = job_dir / "mixed.sdf" - combined.write_text( - good1.read_text() + "this is not a molecule\n$$$$\n" + good2.read_text() + # A corrupt record that ends cleanly at its own $$$$, so the supplier + # yields exactly [ethanol, None, propanol]. Verified: the counts line + # promises two atoms and the atom block does not deliver them. + corrupt = ( + "corrupt\n" + " RDKit 3D\n" + "\n" + " 2 1 0 0 0 0 0 0 0 0999 V2000\n" + "NOT_A_COORD_LINE\n" + " 1 2 1 0\n" + "M END\n" + "$$$$\n" ) + combined = job_dir / "mixed.sdf" + combined.write_text(good1.read_text() + corrupt + good2.read_text()) out = calc_thermo(str(combined), "AIMNET", use_gpu=False)