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..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,17 +60,39 @@ 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) 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 +100,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 +136,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(