diff --git a/CHANGELOG.md b/CHANGELOG.md index c648faae..643af971 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -76,6 +76,58 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 module that owns them. Documented public names are unaffected and still resolve lazily on first access. +- **One model contract, and `batch_opt` no longer imports `model_factory`.** + `optimizing` used to resolve a model name itself; it now takes a built adapter, + and everything after `out_f` is keyword-only. + + | before | after | + |---|---| + | `optimizing(in_f, out_f, "AIMNET", device, config)` | `optimizing(in_f, out_f, adapter=…, device=…, config=…)` | + | `pad_from_mols(mols, model_name, device, coord_pad, species_pad)` | `pad_from_mols(mols, adapter, device)` | + | `from Auto3D.batch_opt.species import …` | `from Auto3D.models.species import …` | + | `to_model_species(nums, "ANI2xt")` | `to_ani2xt_species(nums)` | + | `Auto3D.models.adapter.ModelAdapter` | `Auto3D.models.contract.ModelAdapter` | + + Build the adapter **inside** the worker process — an adapter must not cross a + `spawn` boundary. `ModelAdapter` loses `device` and gains `to_species` and + `energy`; `Auto3D.model_factory.BaseModelAdapter` is no longer re-exported. + `Auto3D.models` still re-exports `ModelAdapter`. + + `CustomNNP` — the *public* custom-NNP contract — is unchanged, and remains + `forward(species, coords, charges) -> energies`. It does lose + `@runtime_checkable`: `isinstance` against a Protocol checks only attribute + presence, so it could not see `Module.forward`'s stub and never told you + anything the validator wasn't already checking by hand. `isinstance(x, CustomNNP)` + now raises `TypeError` and points at `validate_custom_nnp`. + +- **`opt_steps` below 10 is now refused at construction.** `FIELD_BOUNDS` + declared a minimum of 1 while `utils/validation.py` hand-wrote `>= 10` in two + other places — two different minimums for one option. 10 is the correct one: + `n_steps` only tests all-converged on `istep % 10 == 0`, emits progress on the + same cadence, and guards its statistics with an explicit `n >= 10`, so below 10 + there is no early exit, no progress and no reporting. FIRE also needs several + steps to build velocity. `opt_steps < 10` was returning an unconverged + structure labelled as optimized; loosening the bound to 1 would have accepted + that, so the stricter value wins. + +- **`check_valid_configuration` takes an `Auto3DOptions`** instead of ten keyword + arguments. It carried a third set of defaults, including a literal + `opt_steps=2000`, which is exactly how a schema drifts from the one users + configure. `Auto3DOptions` is now the single source of truth, and engine + choices live in one `ENGINE_CHOICES` table. + + Also: `tauto_engine` is validated unconditionally rather than only when + `enumerate_tautomer` is set (the CLI schema already did this, so the two + disagreed), and `check_input` no longer validates `opt_steps`. + +- **`auto3d .yaml` now exits 2, not 1, on a malformed config file.** An + empty file, a non-mapping top level, or a YAML syntax error raised through the + generic handler as exit 1 "Unexpected Error", while the same file through + `auto3d run -c` gave exit 2 `ConfigurationError`. The legacy path now uses the + same loader, so a script gating on exit 2 gets the same answer from both. The + startup banner also moved after validation, so an unrunnable config is no + longer announced as running. + - **One exit-code scheme, used by every command.** `cli/errors.py` has mapped exception types to differentiated exit codes since 3.x -- 0 success, 1 generic, 2 configuration/input, 3 dependency, 4 GPU, 5 model, plus 6 for a diff --git a/docs/source/migration-4.0.rst b/docs/source/migration-4.0.rst index c111b0dc..0055a5f4 100644 --- a/docs/source/migration-4.0.rst +++ b/docs/source/migration-4.0.rst @@ -576,7 +576,7 @@ API changes coords, species, charges = pad_from_mols(mols, model_name, device) # 4.0 - coords, species, charges, atom_mask = pad_from_mols(mols, model_name, device) + coords, species, charges, atom_mask = pad_from_mols(mols, adapter, device) ``atom_mask`` is ``(batch, max_atoms)`` bool, ``True`` for real atoms. Use it instead of comparing species against a padding sentinel. @@ -674,8 +674,8 @@ Species conversion moved index = getidx(atomic_number, model="ANI2xt") # 4.0 - from Auto3D.batch_opt.species import to_model_species, ANI2XT_INDEX - indices = to_model_species(atomic_numbers, "ANI2xt") # whole molecule at once + from Auto3D.models.species import to_ani2xt_species, ANI2XT_INDEX + indices = to_ani2xt_species(atomic_numbers) # whole molecule at once ``energy_tol`` and ``energy_patience`` removed ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -726,16 +726,31 @@ tolerance) is a different parameter and is unchanged. ``AUTO3D_USE_ENSEMBLE`` is no longer read. Passing either argument now raises ``TypeError``, which is the point: misspellings were previously swallowed. -.. warning:: +``optimizing`` takes an adapter, keyword-only +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code:: python - ``optimizing.__init__`` also dropped ``use_ensemble`` from its parameter - list, which shifts ``progress_cb`` into positional slot 6. A legacy - positional caller such as ``optimizing(in_f, out_f, name, device, config, - True)`` now silently binds ``True`` to ``progress_cb`` instead of raising - an error -- ``n_steps`` wraps every progress callback in ``except - Exception: pass``, so a wrong-typed callback is swallowed rather than - surfaced. Call ``optimizing`` with keyword arguments, especially for - ``progress_cb``, to avoid this. + # 3.x + opt = optimizing(in_f, out_f, "AIMNET", device, config) + + # 4.0 + from Auto3D.model_factory import create_model + opt = optimizing(in_f, out_f, adapter=create_model("AIMNET", device), + device=device, config=config) + +``optimizing`` no longer resolves a model name itself -- ``batch_opt`` does not +import ``model_factory`` at all now -- so the caller builds the adapter and +passes it in. Build it inside the worker process: an adapter must not cross a +``spawn`` boundary. + +Everything after ``out_f`` is keyword-only. That closes a hazard an earlier +draft of this guide warned about: when ``use_ensemble`` was dropped from the +parameter list, ``progress_cb`` moved into positional slot 6, so a legacy +positional call silently bound a stray positional argument to ``progress_cb`` +-- and ``n_steps`` wraps every progress callback in ``except Exception: pass``, +so the wrong-typed callback was swallowed rather than surfaced. A positional +call now raises ``TypeError`` immediately. ``Calculator`` and ``mol2aimnet_input`` require ``model_name`` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/parameters.yaml b/parameters.yaml index e438c68e..b1d6f4d3 100644 --- a/parameters.yaml +++ b/parameters.yaml @@ -34,6 +34,14 @@ convergence_threshold: 0.01 patience: 250 batchsize_atoms: 1024 +# Parallel conformer embedding +# Off by default: enabling it spawns worker processes, which changes a run's +# resource profile. Runs stay serial below parallel_embedding_threshold +# molecules even when it is on, since spawning costs more than it saves there. +use_parallel_embedding: False +parallel_workers: 4 +parallel_embedding_threshold: 10 + # Performance settings allow_tf32: False diff --git a/src/Auto3D/ASE/geometry.py b/src/Auto3D/ASE/geometry.py index 41b78a21..82162b4e 100644 --- a/src/Auto3D/ASE/geometry.py +++ b/src/Auto3D/ASE/geometry.py @@ -17,7 +17,7 @@ DEFAULT_CONVERGENCE_THRESHOLD, DEFAULT_OPT_STEPS, ) -from Auto3D.model_factory import get_device +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.energy import E_TOT_HARTREE_PROP, E_TOT_PROP @@ -262,7 +262,13 @@ def opt_geometry( patience=patience if patience is not None else opt_steps, batchsize_atoms=batchsize_atoms, ) - opt_engine = optimizing(path, outpath, model_name, device, opt_config) + # Built here, in the process that runs the optimization. `optimizing` no + # longer constructs its own adapter (audit M41); see the note at + # `Auto3D.workflow_workers.optim_rank_wrapper` about why construction must + # not be hoisted past the frame that does the work. + adapter = create_model(model_name, device) + opt_engine = optimizing(path, outpath, adapter=adapter, device=device, + config=opt_config) opt_engine.run() # `optimizing.run()` already wrote E_tot in Hartree; this pass only adds diff --git a/src/Auto3D/ASE/thermo.py b/src/Auto3D/ASE/thermo.py index 0dda7f94..f189a890 100644 --- a/src/Auto3D/ASE/thermo.py +++ b/src/Auto3D/ASE/thermo.py @@ -24,7 +24,6 @@ from tqdm import tqdm from Auto3D.batch_opt.batchopt import EnForce_ANI -from Auto3D.batch_opt.species import to_model_species from Auto3D.constants import ( DEFAULT_OPT_STEPS, DEFAULT_THERMO_CONVERGENCE_THRESHOLD, @@ -38,6 +37,7 @@ ) from Auto3D.model_factory import create_model, get_device 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 import hartree2ev from Auto3D.utils.logging_config import get_logger diff --git a/src/Auto3D/SPE.py b/src/Auto3D/SPE.py index f9ce69a9..fe5b9b31 100644 --- a/src/Auto3D/SPE.py +++ b/src/Auto3D/SPE.py @@ -152,9 +152,12 @@ def calc_spe( # batch (AIMNet2) must be told which slots are real rather than inferring # it from `species == species_pad`, which deletes a legitimate atomic # number 0 (an R-group `*` atom) along with the padding (audit C13). + # One argument, one source: the adapter supplies the species convention AND + # both pad values. This call used to hand over `model_name` alongside the + # adapter's two pads, so the remap and the sentinel came from different + # places and could contradict each other (audit C3/C4). coord_padded, numbers_padded, charges, atom_mask = pad_from_mols( - mols, model_name, device, - coord_pad=model_adapter.coord_pad, species_pad=model_adapter.species_pad + mols, model_adapter, device ) es, fs = model.forward_batched( diff --git a/src/Auto3D/auto3D.py b/src/Auto3D/auto3D.py index 8442ca33..079d841b 100644 --- a/src/Auto3D/auto3D.py +++ b/src/Auto3D/auto3D.py @@ -22,6 +22,7 @@ from Auto3D.config import Auto3DOptions from Auto3D.exceptions import ConfigurationError from Auto3D.isomers import IsomerEngineFactory +from Auto3D.model_factory import create_model from Auto3D.models.preflight import preflight_model from Auto3D.ranking import ranking from Auto3D.utils import ( @@ -172,18 +173,7 @@ def smiles2mols(smiles: list[str], args: Auto3DOptions) -> list[Chem.Mol]: # gpu_idx) the same way main() does via WorkflowOrchestrator -- # check_input alone does not catch this, so it used to only surface # opaquely deep inside optimization. - config_errors = check_valid_configuration( - path=args.path, - k=args.k, - window=args.window, - use_gpu=args.use_gpu, - gpu_idx=args.gpu_idx, - optimizing_engine=args.optimizing_engine, - isomer_engine=args.isomer_engine, - opt_steps=args.opt_steps, - enumerate_tautomer=args.enumerate_tautomer, - tauto_engine=args.tauto_engine, - ) + config_errors = check_valid_configuration(args) if config_errors: raise ConfigurationError( "Invalid configuration:\n - " + "\n - ".join(config_errors) @@ -228,8 +218,14 @@ def smiles2mols(smiles: list[str], args: Auto3DOptions) -> list[Chem.Mol]: else: device = torch.device("cpu") 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 + # boundary here, but see `Auto3D.workflow_workers.optim_rank_wrapper` + # for why construction must never be hoisted past the frame that works. + adapter = create_model(args.optimizing_engine, device) opt_engine = optimizing(meta["enumerated_sdf"], meta["optimized_og"], - args.optimizing_engine, device, opt_config) + adapter=adapter, device=device, + config=opt_config) opt_engine.run() # Ranking step diff --git a/src/Auto3D/auto3Dcli.py b/src/Auto3D/auto3Dcli.py index dafa44d3..441d7d21 100644 --- a/src/Auto3D/auto3Dcli.py +++ b/src/Auto3D/auto3Dcli.py @@ -86,20 +86,18 @@ def _run_legacy_yaml(yaml_path: str) -> None: # `verbose` key already doubles as the logging-verbosity switch below # (configure_logging); reuse that same key as a coarse opt-in for a # traceback on failure too, rather than always/never showing one. - # `parameters` stays None until (and unless) the YAML actually loads, so - # a failure before that point (bad path, unparsable YAML) falls back to - # no traceback instead of raising a secondary NameError here. + # `verbose` stays 0 until (and unless) the configuration actually + # validates, so a failure before that point (bad path, unparsable YAML) + # falls back to no traceback instead of raising a secondary NameError here. # # `job_hint` follows the same rule for the Ctrl-C report below: None until # there is a configuration to derive a job directory from. - parameters: dict | None = None + verbose: int = 0 job_hint: str | None = None try: - import yaml - from Auto3D.auto3D import main from Auto3D.cli.commands.run import _exit_if_incomplete - from Auto3D.cli.config_schema import build_cli_config + from Auto3D.cli.config_schema import load_yaml_config from Auto3D.cli.results import ( FailedMolecule, WorkflowResults, @@ -112,42 +110,54 @@ def _run_legacy_yaml(yaml_path: str) -> None: if not Path(yaml_path).is_file(): raise InputValidationError(f"Config file not found: {yaml_path}") - with open(yaml_path) as f: - parameters = yaml.safe_load(f) - - # Convert 'None' strings to None - for key, val in list(parameters.items()): - if val == "None": - parameters[key] = None - - configure_logging(verbose=parameters.get("verbose", False)) - - # Print banner - gpu_info = f"CUDA:{parameters.get('gpu_idx', 0)}" if parameters.get("use_gpu", True) else "CPU" - k = parameters.get("k") - window = parameters.get("window") - output_info = f"k={k}" if k else f"window={window}" if window else "k=1" + # `load_yaml_config` is THE YAML ingestion path -- the same function + # `auto3d run -c` calls (cli/commands/run.py). This entry point used to + # carry its own `yaml.safe_load` plus its own "None"-string loop, which + # shared every *value* validator with the modern path (both ended in + # `build_cli_config`) but none of the three *shape* guards: an empty + # file, a non-mapping top level, or a YAML syntax error reached + # `parameters.items()` here and surfaced as AttributeError/TypeError/ + # yaml.YAMLError under the generic "Unexpected Error" panel at exit 1, + # while the identical file through `-c` gave a ConfigurationError at + # exit 2 with a hint. Two exit codes for one file is precisely what + # build_cli_config's docstring says it exists to prevent, so the + # duplicate ingestion is gone rather than patched -- see + # tests/test_legacy_yaml_parity.py, which asserts both entry points + # report the same exception class and the same exit code for each shape. + config = load_yaml_config(Path(yaml_path)) + + # Logging and the banner are derived from the VALIDATED config, not + # from the raw dict. They used to run before validation, reading + # `parameters.get(...)` with its own per-key defaults ("AIMNET", + # gpu_idx 0, use_gpu True) -- a fourth place option defaults were + # written, and one that could print a banner for a configuration about + # to be rejected. Nothing logs between the old and new positions, so no + # log line is lost; the error panel comes from `handle_error`, not from + # logging, so a pre-validation failure still reports normally. + verbose = 1 if config.verbose else 0 + configure_logging(verbose=config.verbose) + + # Before the banner: a settings-only config file (valid for + # `auto3d run INPUT -c`, unrunnable here -- this form has no other + # source of an input path) is refused by `to_auto3d_options` with a + # ConfigurationError naming the missing key, and announcing a run that + # cannot start would be worse than not announcing it. + options = config.to_auto3d_options() + # Same three expressions `cli/commands/run.py` uses, on the same + # validated object. The old `output_info` had a third branch -- + # `else "k=1"` -- advertising a default no entry point applies any + # more: every one of them now refuses a config with neither selector. + gpu_info = f"CUDA:{config.gpu_idx}" if config.use_gpu else "CPU" + output_info = f"k={config.k}" if config.k else f"window={config.window}" print_banner( - input_path=parameters.get("path", "?"), - engine=parameters.get("optimizing_engine", "AIMNET"), + input_path=str(config.path), + engine=config.optimizing_engine, gpu_info=gpu_info, output_info=output_info, ) console.print() - # CLIConfig gives this legacy path the same validation as `auto3d run - # -c`: every Field bound (shared with Auto3DOptions via - # check_field_bounds/FIELD_BOUNDS), the engine registry check, - # parse_gpu_idx, and Literal validation on tauto_engine/isomer_engine. - # It also means extra="forbid": a YAML key CLIConfig doesn't - # recognize now raises -- via build_cli_config, which translates - # pydantic's ValidationError into Auto3D's own ConfigurationError, so - # the blanket `except Exception` below shows exit code 2 with a hint - # instead of the generic "Unexpected Error" panel at exit 1 -- instead - # of silently passing through to Auto3DOptions as it used to. - config = build_cli_config(**parameters) - options = config.to_auto3d_options() job_hint = job_directory_hint(options.path, options.job_name) result = main(options) @@ -194,7 +204,6 @@ def _run_legacy_yaml(yaml_path: str) -> None: # Ctrl-C on this path printed nothing whatsoever. handle_interrupt(job_hint=job_hint, elapsed_seconds=time.time() - start_time) except Exception as e: # noqa: BLE001 - present every failure as a clean panel - verbose = 1 if isinstance(parameters, dict) and parameters.get("verbose") else 0 handle_error(e, verbose=verbose) diff --git a/src/Auto3D/batch_opt/ANI2xt_no_rep.py b/src/Auto3D/batch_opt/ANI2xt_no_rep.py index 6fcebe71..a5cc0836 100644 --- a/src/Auto3D/batch_opt/ANI2xt_no_rep.py +++ b/src/Auto3D/batch_opt/ANI2xt_no_rep.py @@ -3,7 +3,7 @@ import torch import torch.nn as nn -from Auto3D.batch_opt.species import ANI2XT_INDEX +from Auto3D.models.species import ANI2XT_INDEX from Auto3D.utils import hartree2ev # Note: Do NOT set torch.manual_seed() at module level. @@ -17,6 +17,59 @@ root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) ani_2xt_dict = os.path.join(root, "models/ani2xt_no_repulsion.pt") +#: Hidden-layer widths of each per-element network, in ANI2xt's ModuleList +#: order: H, C, N, O, F, S, Cl -- the same order as +#: :data:`Auto3D.models.species.ANI2XT_INDEX`. Every network is +#: ``aev_dim -> w0 -> w1 -> w2 -> 1`` with CELU(0.1) between the layers, so the +#: three widths are the only thing that differs between them. +#: +#: This replaces seven copy-pasted ``nn.Sequential`` blocks (69 lines) that +#: differed only in these integers -- and in which the SOURCE declared S before +#: F while the ``ModuleList`` placed F before S. That mismatch was harmless only +#: because F, S and Cl share widths, and it was invisible without cross-checking +#: two distant lines. The table is now the single statement of the order. +#: +#: Changing the order here silently rewires which element uses which weights: +#: ``nn.ModuleList.load_state_dict`` matches by POSITION (``"0.0.weight"``, +#: ``"1.0.weight"``, ...), never by any Python name. That is also why collapsing +#: the seven blocks into this table did not break the shipped checkpoint -- +#: pinned by ``tests/test_model_adapter.py::TestAni2xtNetworksAreTableDriven``, +#: which loads ``models/ani2xt_no_repulsion.pt`` both ways and compares every +#: tensor. +WIDTHS: tuple[tuple[int, int, int], ...] = ( + (256, 192, 160), # H + (224, 192, 160), # C + (192, 160, 128), # N + (192, 160, 128), # O + (160, 128, 96), # F + (160, 128, 96), # S + (160, 128, 96), # Cl +) + + +def _atomic_mlp(aev_dim: int, widths: tuple[int, int, int]) -> nn.Sequential: + """Build one per-element energy network. + + Args: + aev_dim: Width of the AEV feature vector produced by the AEV computer. + widths: The three hidden-layer widths, as one row of :data:`WIDTHS`. + + Returns: + ``Linear -> CELU -> Linear -> CELU -> Linear -> CELU -> Linear(->1)``, + layer for layer identical to the hand-written blocks this replaced. + """ + w0, w1, w2 = widths + return nn.Sequential( + nn.Linear(aev_dim, w0), + nn.CELU(0.1), + nn.Linear(w0, w1), + nn.CELU(0.1), + nn.Linear(w1, w2), + nn.CELU(0.1), + nn.Linear(w2, 1), + ) + + class ANI2xt(nn.Module): def __init__(self, device, state_dict=ani_2xt_dict, periodic_table_index=False): super().__init__() @@ -44,80 +97,14 @@ def __init__(self, device, state_dict=ani_2xt_dict, periodic_table_index=False): aev_computer = torchani.AEVComputer(radial, angular, num_species) aev_dim = aev_computer.out_dim - H_network = torch.nn.Sequential( - torch.nn.Linear(aev_dim, 256), - torch.nn.CELU(0.1), - torch.nn.Linear(256, 192), - torch.nn.CELU(0.1), - torch.nn.Linear(192, 160), - torch.nn.CELU(0.1), - torch.nn.Linear(160, 1) - ) - - C_network = torch.nn.Sequential( - torch.nn.Linear(aev_dim, 224), - torch.nn.CELU(0.1), - torch.nn.Linear(224, 192), - torch.nn.CELU(0.1), - torch.nn.Linear(192, 160), - torch.nn.CELU(0.1), - torch.nn.Linear(160, 1) - ) - - N_network = torch.nn.Sequential( - torch.nn.Linear(aev_dim, 192), - torch.nn.CELU(0.1), - torch.nn.Linear(192, 160), - torch.nn.CELU(0.1), - torch.nn.Linear(160, 128), - torch.nn.CELU(0.1), - torch.nn.Linear(128, 1) - ) - O_network = torch.nn.Sequential( - torch.nn.Linear(aev_dim, 192), - torch.nn.CELU(0.1), - torch.nn.Linear(192, 160), - torch.nn.CELU(0.1), - torch.nn.Linear(160, 128), - torch.nn.CELU(0.1), - torch.nn.Linear(128, 1) + # One factory + WIDTHS (module level), in place of seven copy-pasted + # nn.Sequential blocks. Order is ModuleList order -- H, C, N, O, F, S, Cl + # -- matching Auto3D.models.species.ANI2XT_INDEX, and it is load-bearing: + # the checkpoint's keys are positional indices. + self.networks = torch.nn.ModuleList( + [_atomic_mlp(aev_dim, widths) for widths in WIDTHS] ) - - S_network = torch.nn.Sequential( - torch.nn.Linear(aev_dim, 160), - torch.nn.CELU(0.1), - torch.nn.Linear(160, 128), - torch.nn.CELU(0.1), - torch.nn.Linear(128, 96), - torch.nn.CELU(0.1), - torch.nn.Linear(96, 1) - ) - - F_network = torch.nn.Sequential( - torch.nn.Linear(aev_dim, 160), - torch.nn.CELU(0.1), - torch.nn.Linear(160, 128), - torch.nn.CELU(0.1), - torch.nn.Linear(128, 96), - torch.nn.CELU(0.1), - torch.nn.Linear(96, 1) - ) - - Cl_network = torch.nn.Sequential( - torch.nn.Linear(aev_dim, 160), - torch.nn.CELU(0.1), - torch.nn.Linear(160, 128), - torch.nn.CELU(0.1), - torch.nn.Linear(128, 96), - torch.nn.CELU(0.1), - torch.nn.Linear(96, 1) - ) - - # Create a ModuleList to hold networks (indexed by species: H=0, C=1, N=2, O=3, F=4, S=5, Cl=6) - self.networks = torch.nn.ModuleList([ - H_network, C_network, N_network, O_network, F_network, S_network, Cl_network - ]) checkpoint = torch.load(state_dict, map_location=device, weights_only=True) self.networks.load_state_dict(checkpoint) # Move networks to device diff --git a/src/Auto3D/batch_opt/batchopt.py b/src/Auto3D/batch_opt/batchopt.py index 6aae893f..f6c7cf49 100644 --- a/src/Auto3D/batch_opt/batchopt.py +++ b/src/Auto3D/batch_opt/batchopt.py @@ -11,22 +11,14 @@ from collections.abc import Callable from Auto3D.config import OptimizationConfig + from Auto3D.models.contract import ModelAdapter logger = get_logger(__name__) -try: - import torchani # noqa: F401 (optional dependency probe) -except ImportError: - pass from collections import defaultdict from rdkit import Chem -try: - from .ANI2xt_no_rep import ANI2xt # noqa: F401 (optional dependency probe) -except ImportError: - pass - # Note: TF32 settings are now configured via Auto3D.torch_config.configure_torch() # and the allow_tf32 option in Auto3DOptions. The hardcoded settings have been # removed to allow user configuration. @@ -39,7 +31,11 @@ # Re-export for backward compatibility (print_stats kept as public re-export) from Auto3D.batch_opt.optimization_engine import n_steps, print_stats # noqa: F401 from Auto3D.constants import INITIAL_ENERGY_SENTINEL, INITIAL_FMAX_SENTINEL -from Auto3D.model_factory import create_model + +# Deliberately NOT `from Auto3D.model_factory import create_model`. This module +# is the numerical layer; the factory sits above it, and importing upward made +# `optimizing` construct its own dependency. Callers inject a ready adapter -- +# see `optimizing.__init__`. from Auto3D.utils.convergence import set_converged from Auto3D.utils.energy import set_e_tot_from_ev from Auto3D.utils.stereo_check import apply_optimized_coords @@ -148,7 +144,8 @@ def __init__( self, in_f: str, out_f: str, - name: str, + *, + adapter: "ModelAdapter", device: torch.device, config: "OptimizationConfig | dict", progress_cb: "Callable[[dict], None] | None" = None, @@ -158,13 +155,30 @@ def __init__( Args: in_f: Input SDF file path. out_f: Output SDF file path. - name: Model name ('AIMNET', 'ANI2x', 'ANI2xt', or path to custom model). + adapter: A ready model adapter satisfying + :class:`Auto3D.models.contract.ModelAdapter`, built by the + caller. This used to be a model NAME that this class handed to + ``Auto3D.model_factory.create_model`` itself -- the numerical + layer importing upward into the construction layer and building + its own dependency (audit M41). + + **The caller must construct it inside the process that will run + the optimization.** See the comments at the two production call + sites (``workflow_workers.optim_rank_wrapper`` and + ``ASE.geometry.opt_geometry``): hoisting construction any further + out pushes a device-resident ``nn.Module`` -- and for AIMNET a + live ``AIMNet2Calculator`` -- across a ``spawn`` boundary. device: Torch device for computation. config: OptimizationConfig dataclass or legacy dict with parameters. + progress_cb: Optional per-step progress callback. + + Everything after ``out_f`` is keyword-only on purpose: the third + positional slot used to be the engine name, and a stale positional caller + would otherwise bind a string silently into the slot that now supplies the + padding values. """ self.in_f = in_f self.out_f = out_f - self.name = name self.device = device self.progress_cb = progress_cb @@ -175,10 +189,12 @@ def __init__( # It's an OptimizationConfig - convert to dict for internal use self._config_dict = config.to_dict() - # Use ModelFactory to create the model adapter - self.model = create_model(name, device) - self.coord_pad = self.model.coord_pad - self.species_pad = self.model.species_pad + # No engine name is retained: `pad_from_mols` asks the adapter for the + # species convention as well as both pad values, so there is nothing left + # for a name to decide here. + self.model = adapter + self.coord_pad = adapter.coord_pad + self.species_pad = adapter.species_pad @property def config(self) -> dict: @@ -246,8 +262,7 @@ def _optimize_bucket(self, bucket_mols, model): position within ``bucket_mols``). """ coord_padded, numbers_padded, charges, atom_mask = pad_from_mols( - bucket_mols, self.name, self.device, - coord_pad=self.coord_pad, species_pad=self.species_pad + bucket_mols, self.model, self.device ) # torch.jit.optimized_execution only affects TorchScript modules; the diff --git a/src/Auto3D/batch_opt/model_wrapper.py b/src/Auto3D/batch_opt/model_wrapper.py index b90f2433..a5513227 100644 --- a/src/Auto3D/batch_opt/model_wrapper.py +++ b/src/Auto3D/batch_opt/model_wrapper.py @@ -6,15 +6,18 @@ """ from __future__ import annotations -from typing import TYPE_CHECKING - import torch import torch.nn as nn from Auto3D.exceptions import OptimizationError -if TYPE_CHECKING: - from Auto3D.model_factory import BaseModelAdapter +# The CONTRACT, not the construction layer. This module used to name +# `Auto3D.model_factory.BaseModelAdapter` -- the numerical layer reaching up into +# the factory for a type that is actually defined below it, and reaching for the +# implementation base class rather than the interface. A runtime (not +# TYPE_CHECKING) import because the gate below consults it; contract.py costs +# nothing beyond torch, which this module already imports. +from Auto3D.models.contract import ModelAdapter, missing_adapter_members class EnForce_ANI(nn.Module): @@ -41,15 +44,22 @@ class EnForce_ANI(nn.Module): def __init__( self, - model_adapter: BaseModelAdapter, + model_adapter: ModelAdapter, batchsize_atoms: int = 1024 * 16, ) -> None: """Initialize EnForce_ANI wrapper. Args: - model_adapter: A model adapter implementing the forward interface. + model_adapter: An object satisfying + :class:`Auto3D.models.contract.ModelAdapter`. Checked here -- + this is the one seam in Auto3D where that Protocol is + load-bearing. batchsize_atoms: Maximum number of atoms per batch (default: 16384). + Raises: + TypeError: ``model_adapter`` is missing contract members, or + ``batchsize_atoms`` is not an int. + The second parameter used to be ``name_or_batchsize: str | int | None``, type-switched between a model name (the pre-adapter API) and a batch size. Passing a string warned that it would be "removed in Auto3D v2.0"; the @@ -57,6 +67,27 @@ def __init__( ever passed one. Removed, so the parameter has one meaning. """ super().__init__() + # What this catches is a CATEGORY ERROR -- a raw nn.Module, a third-party + # calculator, a leftover engine-name string -- named here instead of + # surfacing as an AttributeError several frames deep inside + # forward_batched. It is NOT a contract check: a presence test cannot see + # arity, so an object with a wrong-signature forward still passes, and a + # MagicMock passes trivially (the unit tests rely on that). Do not + # oversell it in this message; the next reader will believe it. + # + # The missing names are computed from the Protocol rather than trusting + # the bare isinstance boolean, so widening ModelAdapter widens this + # message in the same edit. Never use issubclass against ModelAdapter -- + # it raises TypeError for any Protocol with data members. + missing = missing_adapter_members(model_adapter) + if missing: + raise TypeError( + f"EnForce_ANI needs a model adapter satisfying " + f"Auto3D.models.contract.ModelAdapter; " + f"{type(model_adapter).__name__} is missing " + f"{', '.join(missing)}. Build one with " + f"Auto3D.model_factory.create_model." + ) # A caller migrating off the removed API would pass a model name here and, # with the union gone, silently set the batch size to a string -- surfacing # much later inside batching as an unrelated comparison error. Rejected on diff --git a/src/Auto3D/batch_opt/padding.py b/src/Auto3D/batch_opt/padding.py index 5c508105..7fc32b62 100644 --- a/src/Auto3D/batch_opt/padding.py +++ b/src/Auto3D/batch_opt/padding.py @@ -8,15 +8,20 @@ """ from __future__ import annotations +from typing import TYPE_CHECKING + import torch +if TYPE_CHECKING: + # Annotation only, and pointing DOWN the stack: batch_opt depends on + # models/, never the reverse and never on model_factory. + from Auto3D.models.contract import ModelAdapter + def pad_from_mols( mols: list, # List of RDKit Mol objects - model_name: str, + adapter: ModelAdapter, device: torch.device, - coord_pad: float = 0.0, - species_pad: int = -1, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Pad molecular data directly from RDKit Mol objects. @@ -25,11 +30,20 @@ def pad_from_mols( Args: mols: List of RDKit Mol objects with conformers. - model_name: Name of the model ('AIMNET', 'ANI2x', 'ANI2xt', or custom). - Affects how atomic numbers are mapped to species indices. + adapter: The model this batch is being built for, satisfying + :class:`Auto3D.models.contract.ModelAdapter`. It supplies all three + model-dependent pieces -- the species convention + (``adapter.to_species``) and both fill values (``adapter.coord_pad``, + ``adapter.species_pad``). + + This used to be a model-name *string* plus the two pad values as + separate arguments, so both call sites (``SPE.py``, + ``batch_opt/batchopt.py``) handed over a name AND an adapter's pads: + the remap came from one source and the sentinel from another, and + nothing structurally stopped them from contradicting each other. + That is the C3/C4 failure class, and taking the adapter is what makes + it impossible rather than merely absent. device: Target device for tensors (CPU or CUDA). - coord_pad: Padding value for coordinates. Default 0.0. - species_pad: Padding value for species. Default -1. Returns: Tuple of (coords_tensor, species_tensor, charges_tensor, atom_mask) @@ -54,13 +68,13 @@ def pad_from_mols( # Pre-allocate tensors with padding values coords_tensor = torch.full( (batch_size, max_atoms, 3), - coord_pad, + adapter.coord_pad, dtype=torch.float32, device=device ) species_tensor = torch.full( (batch_size, max_atoms), - species_pad, + adapter.species_pad, dtype=torch.long, device=device ) @@ -77,9 +91,7 @@ def pad_from_mols( conf.GetPositions(), dtype=torch.float32, device=device ) - from Auto3D.batch_opt.species import to_model_species - - spec = to_model_species([a.GetAtomicNum() for a in mol.GetAtoms()], model_name) + spec = adapter.to_species([a.GetAtomicNum() for a in mol.GetAtoms()]) species_tensor[i, :n] = torch.tensor(spec, dtype=torch.long, device=device) atom_mask[i, :n] = True diff --git a/src/Auto3D/batch_opt/species.py b/src/Auto3D/batch_opt/species.py deleted file mode 100644 index 5408e761..00000000 --- a/src/Auto3D/batch_opt/species.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Atomic-number to model-species-index conversion. - -This module is the single owner of that mapping. ANI2xt is constructed with -``periodic_table_index=False`` at every site, so its forward expects 0-based -network indices (H=0, C=1, N=2, O=3, F=4, S=5, Cl=6), not atomic numbers. -Every other engine consumes atomic numbers unchanged. - -Before 4.0 this mapping was duplicated in ``utils/chemistry.py`` and -``ASE/thermo.py`` and omitted entirely from the thermo entry points and the -CLI health check, so ANI2xt silently evaluated the wrong species there -(audit findings C3 and C4). -""" -from __future__ import annotations - -from collections.abc import Sequence - -from rdkit import Chem - -from Auto3D.constants import MODEL_ANI2XT - -# Atomic number -> ANI2xt network index. The order matches the ModuleList in -# batch_opt/ANI2xt_no_rep.py; changing one without the other misroutes elements. -ANI2XT_INDEX: dict[int, int] = {1: 0, 6: 1, 7: 2, 8: 3, 9: 4, 16: 5, 17: 6} - -__all__ = ["ANI2XT_INDEX", "to_model_species"] - - -def to_model_species(atomic_numbers: Sequence[int], model_name: str) -> list[int]: - """Convert atomic numbers to the species values a model expects. - - Args: - atomic_numbers: Atomic numbers, one per atom. - model_name: Engine name, matched case-insensitively against - ``"ANI2xt"`` (the same convention ``ModelFactory.create_model`` - uses). Only a match remaps; every other value (AIMNET, any aimnet - registry name, ANI2x, a custom model path) is passed through - unchanged. - - Returns: - Species values in the model's own convention. - - Raises: - ValueError: If ``model_name`` is ``"ANI2xt"`` and an atomic number is - outside its supported set. The message names the atomic number, - the element symbol, and the model. - """ - # Case-insensitive match, matching ModelFactory.create_model's own - # normalization (model_factory.py: name.upper() in cls._adapters). Without - # this, "auto3d models test ani2xt" loads the correct ANI2xt adapter via - # create_model but this function silently passes atomic numbers through - # unconverted -- a C4-shaped bug hiding behind incidental case-matching. - if model_name.upper() != MODEL_ANI2XT.upper(): - return list(atomic_numbers) - - converted: list[int] = [] - for atomic_num in atomic_numbers: - try: - converted.append(ANI2XT_INDEX[atomic_num]) - except KeyError: - symbol = Chem.GetPeriodicTable().GetElementSymbol(atomic_num) - raise ValueError( - f"Element Z={atomic_num} ({symbol}) is not supported by " - f"ANI2xt (supported: H, C, N, O, F, S, Cl)." - ) from None - return converted diff --git a/src/Auto3D/cli/commands/models.py b/src/Auto3D/cli/commands/models.py index e8cfc828..aa83fcbe 100644 --- a/src/Auto3D/cli/commands/models.py +++ b/src/Auto3D/cli/commands/models.py @@ -238,7 +238,6 @@ def execute_models_test( import torch - from Auto3D.batch_opt.species import to_model_species from Auto3D.exceptions import NumericalError from Auto3D.model_factory import create_model, get_device from Auto3D.utils.validation import check_gpu_requested @@ -258,11 +257,13 @@ def execute_models_test( [0.63, -0.63, -0.63], [-0.63, 0.63, -0.63]]], dtype=torch.float, device=device, ) - # Build species in the engine's own convention. Passing raw atomic - # numbers made the ANI2xt check evaluate a Cl+4C species and report - # success (audit C4). + # Build species in the engine's own convention, asked of the + # adapter that will consume them. Passing raw atomic numbers made the + # ANI2xt check evaluate a Cl+4C species and report success (audit + # C4); asking a name-keyed helper instead of the model left the + # convention and the model as two independently-resolved things. species = torch.tensor( - [to_model_species([6, 1, 1, 1, 1], engine)], device=device + [adapter.to_species([6, 1, 1, 1, 1])], device=device ) charges = torch.tensor([0.0], device=device) energy, forces = adapter.forward(coords, species, charges) diff --git a/src/Auto3D/cli/config_schema.py b/src/Auto3D/cli/config_schema.py index 83a1fdb7..fe9545fe 100644 --- a/src/Auto3D/cli/config_schema.py +++ b/src/Auto3D/cli/config_schema.py @@ -7,6 +7,7 @@ from __future__ import annotations +import dataclasses from pathlib import Path from typing import Any, Literal @@ -30,6 +31,66 @@ from Auto3D.exceptions import ConfigurationError from Auto3D.models.preflight import resolve_engine_name +#: ``Auto3DOptions`` fields ``CLIConfig`` deliberately does not carry, with the +#: reason each is excluded. Used by ``to_auto3d_options`` below (which forwards +#: every *other* field mechanically) and by +#: ``tests/test_cli_config_schema.py``'s field-parity test, so "not exposed to +#: the CLI" is stated once instead of once per consumer. +OPTIONS_ONLY_FIELDS: dict[str, str] = { + "input_format": ( + "derived by the workflow from the input file's suffix, not something a " + "user sets" + ), +} + +# Built-in engine names mapped back to the exact spelling Auto3DOptions and the +# model factory compare against. `_validate_engine` accepts any case ( +# `resolve_engine_name` case-folds) but returns the value unchanged, so +# `self.optimizing_engine` still carries whatever the caller typed, e.g. +# "ani2x". Registry names and custom paths are deliberately absent and pass +# through verbatim -- see test_config_accepts_registry_and_path_engines. +_ENGINE_CANONICAL_CASE: dict[str, str] = { + "ANI2X": "ANI2x", + "ANI2XT": "ANI2xt", + "AIMNET": "AIMNET", +} + + +def _to_options_path(value: Path | None) -> str | None: + """``Path`` -> ``str``, keeping an absence an absence. + + ``str(None)`` would be the literal ``"None"``: a path that looks real and + names nothing. + """ + return str(value) if value is not None else None + + +def _to_options_selector(value: Any) -> Any: + """``CLIConfig``'s ``None`` "unset" sentinel -> ``Auto3DOptions``'s ``False``. + + Only ``k``/``window`` need this. The other two ``SENTINEL_FIELDS`` + (``memory``, ``max_confs``) are typed ``int | None`` on *both* classes, so + ``None`` carries across unchanged. + """ + return value if value else False + + +def _to_options_engine(value: str) -> str: + return _ENGINE_CANONICAL_CASE.get(value.upper(), value) + + +#: The only fields whose value differs between the two classes. Everything else +#: is forwarded verbatim by ``to_auto3d_options``, so a field added to both +#: classes cannot be silently dropped on the way across (which the previous +#: hand-written 27-assignment ``return Auto3DOptions(...)`` allowed -- nothing +#: checked that the mapper actually forwarded every field). +_TO_OPTIONS_TRANSFORMS: dict[str, Any] = { + "path": _to_options_path, + "k": _to_options_selector, + "window": _to_options_selector, + "optimizing_engine": _to_options_engine, +} + class CLIConfig(BaseModel): """Validated configuration for Auto3D CLI.""" @@ -65,7 +126,16 @@ class CLIConfig(BaseModel): use_gpu: bool = True gpu_idx: int | list[int] = 0 - # Isomer settings + # Isomer settings. + # + # The two Literals below are the *typed* view of Auto3D.config.ENGINE_CHOICES, + # which is where those whitelists are declared and where + # Auto3DOptions.__post_init__ enforces them. They stay written out here + # because a Literal is what mypy and pydantic's error messages can use and + # neither can read a dict at type-check time -- + # test_engine_choices_table_matches_cliconfig_literals asserts the two agree, + # so this copy cannot drift. (Contrast the numeric bounds, which are in + # FIELD_BOUNDS *only*: a Field(ge=) buys nothing a type can use.) enumerate_tautomer: bool = False tauto_engine: Literal["rdkit", "oechem"] = "rdkit" pKaNorm: bool = True @@ -215,52 +285,24 @@ def to_auto3d_options(self, allow_missing_path: bool = False) -> Auto3DOptions: "-c config.yaml'." ), ) - # Map built-in engine names back to the canonical form expected by - # Auto3DOptions; registry names and custom paths pass through verbatim. - # - # This table is live, not dead: `_validate_engine` (above) now accepts - # any case of these three names -- `resolve_engine_name` case-folds - # them -- but it validates and returns `v` unchanged, so - # `self.optimizing_engine` still carries whatever case the caller - # typed (e.g. "ani2x"). This map is what normalizes that back to the - # exact mixed-case spelling (`ANI2x`/`ANI2xt`) that `MODEL_ANI2X`/ - # `MODEL_ANI2XT` and their downstream exact-match comparisons expect. - # Registry names/aliases are deliberately left out of this map and - # pass through as typed -- see test_config_accepts_registry_and_path_engines. - engine_map = {"ANI2X": "ANI2x", "ANI2XT": "ANI2xt", "AIMNET": "AIMNET"} - engine = engine_map.get(self.optimizing_engine.upper(), self.optimizing_engine) - - return Auto3DOptions( - # `str(None)` would be the literal "None", a path that looks real - # and names nothing; keep the absence an absence. - path=str(self.path) if self.path is not None else None, - k=self.k if self.k else False, - window=self.window if self.window else False, - enumerate_tautomer=self.enumerate_tautomer, - tauto_engine=self.tauto_engine, - pKaNorm=self.pKaNorm, - isomer_engine=self.isomer_engine, - enumerate_isomer=self.enumerate_isomer, - mode_oe=self.mode_oe, - max_confs=self.max_confs, - mpi_np=self.mpi_np, - optimizing_engine=engine, - use_gpu=self.use_gpu, - gpu_idx=self.gpu_idx, - opt_steps=self.opt_steps, - convergence_threshold=self.convergence_threshold, - patience=self.patience, - threshold=self.threshold, - batchsize_atoms=self.batchsize_atoms, - use_parallel_embedding=self.use_parallel_embedding, - parallel_workers=self.parallel_workers, - parallel_embedding_threshold=self.parallel_embedding_threshold, - memory=self.memory, - capacity=self.capacity, - allow_tf32=self.allow_tf32, - verbose=self.verbose, - job_name=self.job_name, - ) + # Driven off `dataclasses.fields(Auto3DOptions)` -- the authoritative + # schema -- rather than 27 hand-written `field=self.field` assignments. + # Those assignments were the one place drift could not be caught: + # `test_cliconfig_covers_all_auto3doptions_fields` compared field *name* + # sets, so deleting a line here silently dropped a user's setting on the + # floor and every test still passed. Iterating the dataclass means a new + # field is forwarded the moment it exists on both classes, and a field + # that exists on only one is a KeyError/TypeError here rather than a + # silent default. + values: dict[str, Any] = {} + for spec in dataclasses.fields(Auto3DOptions): + if spec.name in OPTIONS_ONLY_FIELDS: + continue + value = getattr(self, spec.name) + transform = _TO_OPTIONS_TRANSFORMS.get(spec.name) + values[spec.name] = transform(value) if transform else value + + return Auto3DOptions(**values) def build_cli_config(**kwargs: Any) -> CLIConfig: diff --git a/src/Auto3D/config.py b/src/Auto3D/config.py index 7edb7851..60683018 100644 --- a/src/Auto3D/config.py +++ b/src/Auto3D/config.py @@ -31,7 +31,23 @@ "k": ("ge", 1), "window": ("gt", 0), "mpi_np": ("ge", 1), - "opt_steps": ("ge", 1), + # 10, not 1. This table declared ("ge", 1) while utils/validation.py + # hand-wrote `< 10` twice -- in check_input and again in + # check_valid_configuration -- so one option had two different minimums and + # opt_steps=5 was accepted by Auto3DOptions/CLIConfig, printed a banner, and + # only then failed at run start. 10 is the surviving number because it is + # the one the optimizer is actually built around, not merely the incumbent: + # batch_opt/optimization_engine.py's n_steps checks "have all structures + # converged?" only on `istep % 10 == 0`, emits progress events on the same + # cadence, and guards its stats print with an explicit `n >= 10` to avoid + # `n // 10 == 0`. Below 10 steps none of those fire, so the loop has no + # early exit, no progress, and no reporting -- a regime its own code + # special-cases. Physically, a FIRE run needs several steps just to build + # up velocity and timestep, so fewer than 10 cannot converge a real + # geometry and would hand back an unconverged structure labeled as + # optimized. Loosening to 1 would have accepted exactly that; the two local + # checks are deleted instead and this is now the sole declaration. + "opt_steps": ("ge", 10), "convergence_threshold": ("gt", 0), "patience": ("ge", 1), "threshold": ("gt", 0), @@ -81,6 +97,30 @@ # (merge_configs). SELECTOR_FIELDS: tuple[str, ...] = ("k", "window") +# The permitted values for the two enumerable engine fields, in one table, for +# the same reason FIELD_BOUNDS holds the numeric bounds: the alternative is the +# same whitelist hand-written once per validator, drifting silently. It was +# written three times before this -- CLIConfig's two ``Literal``s +# (cli/config_schema.py) and two local ``valid_isomer_engines`` / +# ``valid_tauto_engines`` sets inside check_valid_configuration +# (utils/validation.py), the latter now deleted. +# +# Checked by Auto3DOptions.__post_init__ below, which already lowercases both +# fields, so the check belongs there rather than in a downstream validator that +# can only be reached by some entry points. CLIConfig keeps its ``Literal``s -- +# a Literal is a *type*, visible to mypy and to pydantic's error messages, not +# a runtime constraint of the kind FIELD_BOUNDS's docstring forbids duplicating +# -- and tests/test_cli_config_schema.py asserts their arguments equal this +# table, so the one remaining hand-written copy cannot drift. +# +# ``mode_oe`` is deliberately absent: its documented values are omega modes +# nothing validates today, and adding a constraint here would be a new +# restriction rather than the consolidation of an existing one. +ENGINE_CHOICES: dict[str, tuple[str, ...]] = { + "isomer_engine": ("rdkit", "omega"), + "tauto_engine": ("rdkit", "oechem"), +} + _BOUND_OPS: dict[str, tuple[object, str]] = { "ge": (operator.ge, ">="), "gt": (operator.gt, ">"), @@ -184,6 +224,37 @@ def check_selectors_mutually_exclusive(values: dict) -> None: ) +def check_engine_choices(values: dict) -> None: + """Validate ``values`` (field name -> value) against ``ENGINE_CHOICES``. + + Fields missing from ``values`` are skipped, so callers may pass a partial + mapping (the same convention ``check_field_bounds`` uses). Values are + compared case-insensitively; ``Auto3DOptions.__post_init__`` has already + lowercased both fields by the time it calls this, so the fold only matters + for a direct caller. + + Unconditional, unlike the check it replaces: ``check_valid_configuration`` + validated ``tauto_engine`` only when ``enumerate_tautomer`` was true, while + ``CLIConfig``'s ``Literal["rdkit", "oechem"]`` has always rejected a bad + value regardless. That was an entry-point divergence -- ``Auto3DOptions( + tauto_engine="bogus")`` was accepted from Python and refused from the CLI -- + so the stricter of the two is what survives. + + Raises: + ConfigurationError: naming the field, the received value, and the + permitted set. + """ + for name, choices in ENGINE_CHOICES.items(): + if name not in values: + continue + value = values[name] + if isinstance(value, str) and value.lower() in choices: + continue + raise ConfigurationError( + f"{name} must be one of {', '.join(choices)}, got {value!r}" + ) + + def optimizer_worker_indices( use_gpu: bool, gpu_idx: "int | list[int]" ) -> list[int]: @@ -344,10 +415,11 @@ class Auto3DOptions: survives dataclasses.replace()/pickling and stays in the dict-like API.""" def __post_init__(self): - """Normalize string values to lowercase and validate ranges.""" + """Normalize string values to lowercase, then validate choices and ranges.""" self.tauto_engine = self.tauto_engine.lower() self.isomer_engine = self.isomer_engine.lower() self.mode_oe = self.mode_oe.lower() + check_engine_choices({name: getattr(self, name) for name in ENGINE_CHOICES}) check_field_bounds({name: getattr(self, name) for name in FIELD_BOUNDS}) def __getitem__(self, key: str): @@ -393,7 +465,7 @@ class OptimizationConfig: Example: >>> config = OptimizationConfig(opt_steps=1000, convergence_threshold=0.005) - >>> optimizer = optimizing(in_f, out_f, model, device, config) + >>> optimizer = optimizing(in_f, out_f, adapter=adapter, device=device, config=config) """ opt_steps: int = DEFAULT_OPT_STEPS diff --git a/src/Auto3D/model_factory.py b/src/Auto3D/model_factory.py index 030e597f..31364fc5 100644 --- a/src/Auto3D/model_factory.py +++ b/src/Auto3D/model_factory.py @@ -3,6 +3,7 @@ import os from pathlib import Path +from typing import TYPE_CHECKING import torch @@ -18,10 +19,19 @@ AIMNet2Adapter, ANI2xAdapter, ANI2xtAdapter, - BaseModelAdapter, CustomModelAdapter, ) +if TYPE_CHECKING: + # Annotation-only. Every signature here promises the CONTRACT + # (Auto3D.models.contract.ModelAdapter), not the implementation base class; + # `_adapters` is the sole exception, because it really is a registry of + # Auto3D's own classes. Kept behind TYPE_CHECKING so + # `Auto3D.model_factory.BaseModelAdapter` is no longer an incidental + # re-export -- import it from Auto3D.models. + from Auto3D.models.adapter import BaseModelAdapter + from Auto3D.models.contract import ModelAdapter + # Environment variable to enable torch.compile() by default _COMPILE_ENV_VAR = "AUTO3D_COMPILE_MODEL" @@ -54,7 +64,7 @@ class ModelFactory: assert set(_adapters) == set(BUILTIN_ANI_MODELS) # Model instance cache: key = (name, device_str, compile_model) - _cache: dict[tuple[str, str, bool], BaseModelAdapter] = {} + _cache: dict[tuple[str, str, bool], ModelAdapter] = {} @classmethod def clear_cache(cls) -> None: @@ -81,7 +91,7 @@ def create( device: torch.device | None = None, compile_model: bool | None = None, use_cache: bool = True, - ) -> BaseModelAdapter: + ) -> ModelAdapter: """Create a model adapter by name. Args: @@ -176,7 +186,7 @@ def create_model( device: torch.device | None = None, compile_model: bool | None = None, use_cache: bool = True, -) -> BaseModelAdapter: +) -> ModelAdapter: """Convenience function to create a model adapter. Args: diff --git a/src/Auto3D/models/__init__.py b/src/Auto3D/models/__init__.py index fcda6977..a620d310 100644 --- a/src/Auto3D/models/__init__.py +++ b/src/Auto3D/models/__init__.py @@ -1,13 +1,23 @@ -"""Model adapters and neural network potentials for Auto3D.""" +"""Model contracts, adapters and neural network potentials for Auto3D. + +Both contracts live in :mod:`Auto3D.models.contract`: + +* :class:`~Auto3D.models.contract.CustomNNP` -- what a user's own NNP must + satisfy, ``forward(species, coords, charges) -> energies``. +* :class:`~Auto3D.models.contract.ModelAdapter` -- what Auto3D's internals talk + to, ``forward(coords, species, charges, atom_mask=None) -> (energies, forces)``. + +The two take ``species`` and ``coords`` in opposite order, deliberately and +permanently. Read that module's docstring before touching either. +""" from Auto3D.models.adapter import ( AIMNet2Adapter, ANI2xAdapter, ANI2xtAdapter, BaseModelAdapter, CustomModelAdapter, - ModelAdapter, ) -from Auto3D.models.contract import CustomNNP +from Auto3D.models.contract import CustomNNP, ModelAdapter __all__ = [ "CustomNNP", diff --git a/src/Auto3D/models/adapter.py b/src/Auto3D/models/adapter.py index a0ef1a1a..755e4bda 100644 --- a/src/Auto3D/models/adapter.py +++ b/src/Auto3D/models/adapter.py @@ -1,10 +1,21 @@ # src/Auto3D/models/adapter.py -"""Model adapters providing consistent interface for all NNP models.""" +"""Implementations of the adapter contract, one per NNP backend. + +The contract itself -- :class:`Auto3D.models.contract.ModelAdapter` -- lives in +:mod:`Auto3D.models.contract`, next to the custom-NNP contract it is so easily +confused with. This module holds only implementations. + +Layering: :mod:`Auto3D.models` is a leaf. It imports ``torch``, +``Auto3D.constants``, ``Auto3D.exceptions`` and its own submodules, and nothing +else from Auto3D. There is exactly one deliberate back-edge into +``Auto3D.batch_opt`` (``ANI2xtAdapter.__init__``'s deferred ``ANI2xt`` import); +see the comment there before moving it. +""" from __future__ import annotations import warnings from abc import ABC, abstractmethod -from typing import Protocol, runtime_checkable +from collections.abc import Sequence import torch import torch.nn as nn @@ -12,6 +23,7 @@ from Auto3D.constants import HARTREE_TO_EV from Auto3D.exceptions import NumericalError from Auto3D.models.loading import load_custom_nnp +from Auto3D.models.species import to_ani2xt_species def _try_compile(model: nn.Module, mode: str = "default") -> nn.Module: @@ -86,55 +98,24 @@ def _validate_outputs(energy: torch.Tensor, forces: torch.Tensor) -> None: ) -@runtime_checkable -class ModelAdapter(Protocol): - """Protocol defining the standard interface for NNP model adapters. - - All model adapters must implement this interface to ensure consistent - behavior across different neural network potential backends. - - This protocol is runtime_checkable, allowing isinstance() checks. - """ - - coord_pad: float - species_pad: int - device: torch.device - - def forward( - self, - coords: torch.Tensor, - species: torch.Tensor, - charges: torch.Tensor, - atom_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Compute energies and forces. - - Args: - coords: Atomic coordinates (batch, n_atoms, 3). - species: Atomic numbers (batch, n_atoms). - charges: Molecular charges (batch,). - atom_mask: Boolean (batch, n_atoms), True for real atoms and False - for padded slots, as returned by - :func:`Auto3D.batch_opt.padding.pad_from_mols`. Required from - any caller that passes a PADDED batch; ``None`` means every - slot holds a real atom. An adapter must never re-derive this by - comparing ``species`` against ``species_pad`` (audit C13). - - Returns: - Tuple of (energies, forces) where energies has shape (batch,) - and forces has shape (batch, n_atoms, 3). Units: eV. - """ - ... - - class BaseModelAdapter(ABC, nn.Module): - """Base class for model adapters. + """Implementation base for Auto3D's adapters. NOT the contract. + + The contract is :class:`Auto3D.models.contract.ModelAdapter`, and that -- not + this class -- is what every signature that wants "an adapter" annotates. + This distinction is the point: production has always accepted structural + implementations (test doubles, and anything a downstream user writes), so + annotating the ABC while accepting the Protocol is exactly what made the + Protocol decorative. The one place this class legitimately appears as a type + is ``ModelFactory._adapters``, a registry of Auto3D's OWN classes. Provides common functionality for all NNP model adapters including: - Model storage and device management - Padding value configuration - Gradient disabling for model parameters (weights are frozen) - Optional torch.compile() for performance optimization + - Concrete ``to_species`` (identity) and ``energy`` defaults, so a subclass + satisfies the contract by implementing ``forward`` alone Note on torch.inference_mode(): This class CANNOT use torch.inference_mode() or torch.no_grad() in forward @@ -187,6 +168,40 @@ def __init__( self.model = model + def to_species(self, atomic_numbers: Sequence[int]) -> list[int]: + """Identity: this model consumes raw atomic numbers. + + Correct for AIMNet2, for ANI2x (constructed with + ``periodic_table_index=True``), and for every custom NNP -- a custom + model declares its own ``species_pad`` and receives atomic numbers, so + remapping them here would silently feed every third-party model + different species indices than its author tested against. ANI2xt is the + sole override; see :meth:`ANI2xtAdapter.to_species`. + """ + return list(atomic_numbers) + + def energy( + self, + coords: torch.Tensor, + species: torch.Tensor, + charges: torch.Tensor, + atom_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Energies only, graph-connected, at the dtype of ``coords``. + + The default takes ``forward``'s first output. That is safe only for + adapters whose ``forward`` is already dtype-preserving and does not + mutate ``coords.requires_grad``; ``ANI2xtAdapter``, ``ANI2xAdapter`` and + ``CustomModelAdapter`` each override it for exactly that reason (the + latter two call ``coords.float()``, which would turn an fp64 caller's + request into an fp32 answer with no error). + + No ``no_grad`` here, deliberately: a caller differentiating this (a + Hessian) needs the graph, and a caller that does not want it wraps its + own call site. + """ + return self.forward(coords, species, charges, atom_mask)[0] + @abstractmethod def forward( self, @@ -266,6 +281,26 @@ def calculator(self): """ return self._calc + def energy( + self, + coords: torch.Tensor, + species: torch.Tensor, + charges: torch.Tensor, + atom_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Energies (eV) via ``forward``; the base default, made explicit. + + Written out rather than inherited so the dtype reasoning is on the record + for the one adapter where it differs. ``forward`` returns float64 + energies whatever it was fed -- an UPCAST, so there is no silent + precision loss to guard against (the hazard the other two overrides + exist for), and whole-graph fp64 through AIMNet2 would be false + precision regardless. Routed through ``forward`` (hence the calculator's + ``forces=True`` path) because that is the route the calculator + guarantees stays connected to ``coord`` in the autograd graph. + """ + return self.forward(coords, species, charges, atom_mask)[0] + def forward( self, coords: torch.Tensor, @@ -344,10 +379,48 @@ def __init__(self, device: torch.device, compile_model: bool = False) -> None: device: Target device for computations. compile_model: Whether to apply torch.compile() for optimization. """ + # THE ONE deliberate back-edge from Auto3D.models into + # Auto3D.batch_opt, and it MUST stay inside this method. Promoting it to + # 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 model = ANI2xt(device) super().__init__(model, device, coord_pad=0.0, species_pad=-1, compile_model=compile_model) + def to_species(self, atomic_numbers: Sequence[int]) -> list[int]: + """Remap atomic numbers to ANI2xt's 0-based network indices. + + ANI2xt is built with ``periodic_table_index=False`` everywhere, so its + ``forward`` expects H=0, C=1, N=2, O=3, F=4, S=5, Cl=6 -- not atomic + numbers. The remap lives on the adapter (rather than in a name-keyed free + function the caller had to remember to invoke) so it cannot be omitted at + one call site and applied at another; that omission is audit findings + C3/C4, where thermo and the CLI health check silently scored a different + molecule than the one submitted. + + Raises: + ValueError: An atomic number outside ANI2xt's element set. + """ + return to_ani2xt_species(atomic_numbers) + + def energy( + self, + coords: torch.Tensor, + species: torch.Tensor, + charges: torch.Tensor, + atom_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Energies (eV) with no ``requires_grad_`` mutation of ``coords``. + + ``forward`` calls ``coords.requires_grad_(True)`` because it must + differentiate to get forces. ``energy`` cannot: an autograd-Hessian + caller hands in a NON-LEAF tensor, and ``requires_grad_`` on a non-leaf + raises. Energies come out float64 (see ``ANI2xt.forward``); coords are + consumed at whatever dtype they arrive in. + """ + return self.model(species, coords) + def forward( self, coords: torch.Tensor, @@ -399,6 +472,27 @@ def __init__(self, device: torch.device, compile_model: bool = False) -> None: model = torchani.models.ANI2x(periodic_table_index=True).to(device) super().__init__(model, device, coord_pad=0.0, species_pad=-1, compile_model=compile_model) + def energy( + self, + coords: torch.Tensor, + species: torch.Tensor, + charges: torch.Tensor, + atom_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Energies (eV) at the dtype of ``coords`` -- NO float32 downcast. + + This override exists solely to keep that promise. ``forward`` calls + ``coords.float()`` for compatibility with torchani's float32 weights; + inheriting ``BaseModelAdapter.energy`` (which is ``forward(...)[0]``) + would therefore turn an fp64 caller's request into an fp32 answer with no + error, no warning, and no way to notice -- the caller that wants fp64 is + computing a Hessian, and it would silently get an fp32 one. Feeding the + model the dtype it was handed pushes that choice back to the caller, + which is the layer that also has to promote the model's weights + (``.double()``) for it to be meaningful. + """ + return self.model((species, coords)).energies * HARTREE_TO_EV + def forward( self, coords: torch.Tensor, @@ -490,6 +584,27 @@ def __init__( model, device, model.coord_pad, model.species_pad, compile_model=False ) + def energy( + self, + coords: torch.Tensor, + species: torch.Tensor, + charges: torch.Tensor, + atom_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Energies (eV) at the dtype of ``coords`` -- NO float32 downcast. + + Same reason as :meth:`ANI2xAdapter.energy`: ``forward`` casts coords and + charges to float32 (documented in the class docstring), so inheriting the + ``forward(...)[0]`` default would silently answer an fp64 request in fp32. + ``charges`` follows ``coords``' dtype so a model that indexes or + concatenates the two does not hit a mismatch. + + The published contract is ``forward(species, coords, charges)`` -- species + FIRST -- and that order is the user's, not this adapter's; see + :class:`Auto3D.models.contract.CustomNNP`. + """ + return self.model(species, coords, charges.to(coords.dtype)) + def forward( self, coords: torch.Tensor, diff --git a/src/Auto3D/models/contract.py b/src/Auto3D/models/contract.py index 102f988c..fcd26a1f 100644 --- a/src/Auto3D/models/contract.py +++ b/src/Auto3D/models/contract.py @@ -1,28 +1,39 @@ # src/Auto3D/models/contract.py -"""The single definition of the custom-NNP contract, next to what enforces it. +"""Both model contracts, in one file, next to what enforces each of them. -Auto3D has two distinct model interfaces, and conflating them is the whole -reason this module exists: +Auto3D has two distinct and *permanently* separate model interfaces. Conflating +them is the whole reason this module exists, and the reason both declarations +live here rather than one file apart: they are forty lines from each other and +cannot be read separately. -* **The custom-NNP contract (this module).** What a *user* implements and hands - to Auto3D as ``optimizing_engine=/path/to/model.pt``. Auto3D calls it as - ``model(species, coords, charges) -> energies`` and differentiates the - returned energy with respect to ``coords`` to obtain forces +* **Contract A -- the custom-NNP contract** (:class:`CustomNNP`). What a *user* + implements and hands to Auto3D as ``optimizing_engine=/path/to/model.pt``. + Auto3D calls it as ``model(species, coords, charges) -> energies`` and + differentiates the returned energy with respect to ``coords`` to obtain forces (:meth:`Auto3D.models.adapter.CustomModelAdapter.forward`). The model returns - energies only; it must not return forces. - -* **The adapter interface** (:class:`Auto3D.models.adapter.ModelAdapter`). - Internal to Auto3D: ``forward(coords, species, charges) -> (energies, - forces)``. Only Auto3D's own adapters implement it. Users never do. - -Note that the two take ``species`` and ``coords`` in *opposite* order. A model -written against the adapter interface silently computes an energy from -transposed tensors and then fails deep inside ``torch.autograd.grad``, so -:func:`validate_custom_nnp` rejects that shape at load time instead. + energies only; it must not return forces. Enforced by + :func:`validate_custom_nnp` at load time. + +* **Contract B -- the adapter interface** (:class:`ModelAdapter`). Internal to + Auto3D: ``forward(coords, species, charges, atom_mask=None) -> (energies, + forces)``. Only Auto3D's own adapters implement it (via + :class:`Auto3D.models.adapter.BaseModelAdapter`, which supplies working + defaults for everything but ``forward``). Users never do. Enforced by + :func:`missing_adapter_members`, which + :class:`Auto3D.batch_opt.model_wrapper.EnForce_ANI` consults on construction. + +Note that the two take ``species`` and ``coords`` in *opposite* order, and that +this is deliberate rather than an accident awaiting cleanup: Contract A's order +is published (``CHANGELOG.md``, ``docs/source/howto/custom_nnp.rst``) and +changing it would break every working third-party model. A model written against +the adapter interface silently computes an energy from transposed tensors and +then fails deep inside ``torch.autograd.grad``, so :func:`validate_custom_nnp` +rejects that shape at load time instead. """ from __future__ import annotations import inspect +from collections.abc import Sequence from typing import Any, Protocol, runtime_checkable import torch @@ -32,10 +43,41 @@ #: Human-readable form of the contract, quoted in every rejection message. EXPECTED_SIGNATURE = "forward(self, species, coords, charges) -> energies" + +def _protocol_data_members(proto: type) -> tuple[str, ...]: + """Names annotated in a Protocol's own class body, in declaration order. + + Data members land in ``__annotations__``; methods do not. Implemented by + hand because the stdlib alternatives are unavailable or unstable here: + ``typing.get_protocol_members`` is 3.13+, and ``__protocol_attrs__`` is a + CPython implementation detail that also folds in the methods. + + ``from __future__ import annotations`` turns the annotation *values* into + strings; the *keys* are unaffected, and only keys are used here. + """ + return tuple(proto.__annotations__) + + +def _protocol_members(proto: type) -> tuple[str, ...]: + """Every member a Protocol declares: data members then public methods.""" + data = _protocol_data_members(proto) + methods = tuple( + name + for name, value in vars(proto).items() + if not name.startswith("_") and callable(value) and name not in data + ) + return data + methods + + #: Attributes every custom NNP must define. They are the padding *fill* values #: Auto3D writes into the batched tensors; ``species_pad`` in particular decides #: what lands in the species tensor's unused slots, so guessing it is unsafe. -REQUIRED_ATTRIBUTES = ("coord_pad", "species_pad") +#: +#: DERIVED from :class:`CustomNNP`'s own annotations rather than retyped beside +#: them, so the two cannot drift. See the warning above ``CustomNNP``: this makes +#: the Protocol's annotations load-bearing. +#: (Assigned below the class, which must exist first.) +REQUIRED_ATTRIBUTES: tuple[str, ...] # Parameter-name vocabulary used only to detect a *transposed* forward. Names # outside these sets are not an error -- they just make the argument order @@ -49,7 +91,25 @@ _CHARGES_NAMES = frozenset({"charges", "charge", "q"}) -@runtime_checkable +# WARNING -- the annotated members of this class are load-bearing, not +# documentation. REQUIRED_ATTRIBUTES is derived from them (below), and +# validate_custom_nnp skips the `forward` signature check entirely for a +# TorchScript RecursiveScriptModule, so for every archive in the wild those +# annotations are the ONLY gate. Adding a third annotated field here instantly +# rejects every existing custom NNP that does not carry it: that is a BREAKING +# change and must be released as one. `test_custom_nnp_contract.py:: +# test_customnnp_data_members_are_exactly_the_two_padding_values` pins the set +# so it cannot happen by accident. +# +# Deliberately NOT @runtime_checkable, unlike ModelAdapter below. A runtime +# Protocol check tests attribute *presence* only, and torch installs +# `Module.forward = _forward_unimplemented` on every nn.Module, so the single +# most common real failure -- a saved module with no forward of its own -- would +# pass isinstance() and then raise NotImplementedError inside the optimization +# loop. `_check_forward_signature` already handles that case correctly and by +# hand, and a bare boolean cannot carry the diagnosis validate_custom_nnp emits. +# So `isinstance(x, CustomNNP)` raises TypeError, which is the honest answer to +# "can I check this at runtime?" -- no; call validate_custom_nnp. class CustomNNP(Protocol): """Protocol a user-supplied NNP must satisfy to be used as an engine. @@ -107,6 +167,142 @@ def forward( ... +REQUIRED_ATTRIBUTES = _protocol_data_members(CustomNNP) + + +@runtime_checkable +class ModelAdapter(Protocol): + """Contract B: the interface Auto3D's own model adapters present. + + Every consumer inside Auto3D -- the optimizer, the single-point-energy path, + the ASE calculator, the CLI health check -- talks to a model through exactly + these members, and this is the type they annotate. Implementations live in + :mod:`Auto3D.models.adapter`; :class:`~Auto3D.models.adapter.BaseModelAdapter` + supplies working defaults for everything except ``forward``, so an in-tree + adapter satisfies this by inheritance. + + Note the argument order is the REVERSE of :class:`CustomNNP` + (``species`` first there, ``coords`` first here) and that this one returns + ``(energies, forces)`` while a custom NNP returns energies only. See the + module docstring. + + ``device`` is deliberately NOT part of this contract. Nothing outside an + adapter reads it (``BaseModelAdapter`` keeps ``self.device`` as an + implementation detail), and requiring it would make every legitimate + structural implementation -- including test doubles that never touch a + device -- non-conforming for no benefit. + + This Protocol IS ``@runtime_checkable``, and unlike before it is actually + consulted: :func:`missing_adapter_members` is called by + :class:`Auto3D.batch_opt.model_wrapper.EnForce_ANI`. Be clear about what that + buys, because overselling it is how the decorator became decorative in the + first place: a presence check catches a **category error** (a raw + ``nn.Module``, a third-party calculator, an engine-name string), NOT a + contract violation. Presence is not arity, and a ``MagicMock`` satisfies it + trivially. Never use ``issubclass`` against this Protocol -- it raises + ``TypeError`` for any Protocol with data members. + """ + + coord_pad: float + """Fill value the batch padder writes into unused coordinate slots.""" + + species_pad: int + """Fill value the batch padder writes into unused species slots.""" + + def to_species(self, atomic_numbers: Sequence[int]) -> list[int]: + """Convert atomic numbers into this model's own species convention. + + The species convention is a property of the *model*, so it lives on the + same object that supplies ``species_pad``. That is what makes it + impossible for the remap and the padding sentinel to come from two + different sources and contradict each other -- the shape of audit + findings C3/C4, where a name-keyed converter and an adapter-supplied pad + disagreed about which slots were padding. + + Args: + atomic_numbers: Atomic numbers, one per atom. + + Returns: + Species values in the model's convention. The identity for every + engine except ANI2xt, which uses 0-based network indices. + """ + ... + + def forward( + self, + coords: torch.Tensor, + species: torch.Tensor, + charges: torch.Tensor, + atom_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute energies and forces. + + Args: + coords: Atomic coordinates (batch, n_atoms, 3). + species: Species in this adapter's own convention (batch, n_atoms), + i.e. the output of :meth:`to_species`. + charges: Molecular charges (batch,). + atom_mask: Boolean (batch, n_atoms), True for real atoms and False + for padded slots, as returned by + :func:`Auto3D.batch_opt.padding.pad_from_mols`. Required from + any caller that passes a PADDED batch; ``None`` means every + slot holds a real atom. An adapter must never re-derive this by + comparing ``species`` against ``species_pad`` (audit C13). + + Returns: + Tuple of (energies, forces) where energies has shape (batch,) + and forces has shape (batch, n_atoms, 3). Units: eV. + """ + ... + + def energy( + self, + coords: torch.Tensor, + species: torch.Tensor, + charges: torch.Tensor, + atom_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Compute energies only, graph-connected and at the caller's dtype. + + Two properties are part of the contract and neither is optional: + + * **No internal ``no_grad``/``inference_mode``.** The result must stay + connected to ``coords`` so a caller can differentiate it (a Hessian). + A caller that wants no graph wraps its own call site. + * **Dtype-preserving.** ``forward`` downcasts to float32 in two adapters + for compatibility with float32 NNP weights. ``energy`` must not: an + fp64 caller that silently receives an fp32 result gets no error and no + warning, only a wrong number. + + Args: + coords: Atomic coordinates (batch, n_atoms, 3). + species: Species in this adapter's own convention (batch, n_atoms). + charges: Molecular charges (batch,). + atom_mask: As for :meth:`forward`. + + Returns: + Energies, shape (batch,), in eV. + """ + ... + + +def missing_adapter_members(obj: Any) -> list[str]: + """Members :class:`ModelAdapter` requires that ``obj`` does not provide. + + Derived from the Protocol, never hand-listed, so widening + :class:`ModelAdapter` widens the rejection message in the same edit. + + Args: + obj: The candidate adapter. + + Returns: + Missing member names in declaration order; empty if ``obj`` structurally + conforms. Remember this is a presence check: an empty list means "not a + category error", not "correct". + """ + return [name for name in _protocol_members(ModelAdapter) if not hasattr(obj, name)] + + def _classify(name: str) -> str | None: """Map a parameter name to ``'species'``/``'coords'``/``'charges'``, or None.""" lowered = name.lower() @@ -162,7 +358,8 @@ def _check_forward_signature(model: Any, source: str) -> None: ) try: - parameters = list(inspect.signature(forward).parameters.values()) + signature = inspect.signature(forward) + parameters = list(signature.parameters.values()) except (ValueError, TypeError): # A TorchScript RecursiveScriptModule's forward is a pybind11 builtin # with no Python signature, so inspect.signature raises ValueError. Skip @@ -181,10 +378,19 @@ def _check_forward_signature(model: Any, source: str) -> None: in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) ] required = [p for p in positional if p.default is inspect.Parameter.empty] + # Render the signature the way the interpreter does, NOT by comma-joining + # parameter names. Joining names drops `*`, `/` and defaults, so a + # keyword-only forward was refused with text identical to the text it was + # being asked for ("has forward(species, coords, charges) ... Expected + # forward(self, species, coords, charges)") -- a message that shows the + # author nothing wrong. The marker IS the explanation. + # The return annotation is dropped so the rendering stays comparable with + # EXPECTED_SIGNATURE's parameter list; everything else -- markers, defaults, + # parameter annotations -- is kept exactly as written. + rendered = str(signature.replace(return_annotation=inspect.Signature.empty)) if len(positional) < 3 or len(required) > 3: - rendered = ", ".join(p.name for p in parameters) or "" raise ModelLoadError( - f"Custom NNP at {source} has forward({rendered}), which cannot be " + f"Custom NNP at {source} has forward{rendered}, which cannot be " f"called with three positional arguments. Expected " f"{EXPECTED_SIGNATURE}, returning energies of shape (batch,) in eV." ) @@ -194,9 +400,8 @@ def _check_forward_signature(model: Any, source: str) -> None: # Unrecognized parameter names: the order is unknowable, so accept. return if classified != ["species", "coords", "charges"]: - rendered = ", ".join(p.name for p in positional[:3]) raise ModelLoadError( - f"Custom NNP at {source} has forward({rendered}), but Auto3D calls a " + f"Custom NNP at {source} has forward{rendered}, but Auto3D calls a " f"custom NNP as {EXPECTED_SIGNATURE} -- species first, coords " f"second. Note this is the opposite order from Auto3D's internal " f"ModelAdapter interface (coords, species, charges), which returns " diff --git a/src/Auto3D/models/species.py b/src/Auto3D/models/species.py new file mode 100644 index 00000000..de0720de --- /dev/null +++ b/src/Auto3D/models/species.py @@ -0,0 +1,100 @@ +# src/Auto3D/models/species.py +"""ANI2xt's species convention, owned by ``models/`` because it is the model's. + +ANI2xt is constructed with ``periodic_table_index=False`` at every site, so its +forward expects 0-based network indices (H=0, C=1, N=2, O=3, F=4, S=5, Cl=6), not +atomic numbers. Every other engine consumes atomic numbers unchanged. + +This table used to live in ``batch_opt/species.py``, which made ``batch_opt`` a +shared-utility host for ``ASE/``, ``cli/`` and ``models/``'s own padder -- three +layers with no business depending on the optimizer package. The species +convention is a property of *the model*, so the model layer owns it, and the +canonical way to reach it is +:meth:`Auto3D.models.contract.ModelAdapter.to_species` -- asking the object that +also supplies ``species_pad``, so the remap and the padding sentinel cannot come +from two sources and disagree (audit findings C3/C4). + +Layering note: rdkit is imported lazily, inside the error branch that needs it +for an element symbol. ``models/`` is reached from ``utils/validation.py`` and +therefore from ``import Auto3D.utils``; a module-scope rdkit import here would +make that import pay for rdkit (audit M44). +""" +from __future__ import annotations + +from collections.abc import Sequence + +from Auto3D.constants import MODEL_ANI2XT + +# Atomic number -> ANI2xt network index. The order matches the ModuleList in +# batch_opt/ANI2xt_no_rep.py; changing one without the other misroutes elements. +ANI2XT_INDEX: dict[int, int] = {1: 0, 6: 1, 7: 2, 8: 3, 9: 4, 16: 5, 17: 6} + +__all__ = ["ANI2XT_INDEX", "to_ani2xt_species", "to_model_species"] + + +def to_ani2xt_species(atomic_numbers: Sequence[int]) -> list[int]: + """Convert atomic numbers to ANI2xt's 0-based network indices. + + The single implementation of this remap. ``ANI2xtAdapter.to_species`` + delegates here, and so does :func:`to_model_species`. + + Args: + atomic_numbers: Atomic numbers, one per atom. + + Returns: + ANI2xt network indices in the same order. + + Raises: + ValueError: An atomic number outside ANI2xt's supported set. The message + names the atomic number, the element symbol, and the model. + """ + converted: list[int] = [] + for atomic_num in atomic_numbers: + try: + converted.append(ANI2XT_INDEX[int(atomic_num)]) + except KeyError: + # Deferred so the happy path never imports rdkit (see module + # docstring); only the message needs the element symbol. + from rdkit import Chem + + symbol = Chem.GetPeriodicTable().GetElementSymbol(int(atomic_num)) + raise ValueError( + f"Element Z={atomic_num} ({symbol}) is not supported by " + f"ANI2xt (supported: H, C, N, O, F, S, Cl)." + ) from None + return converted + + +def to_model_species(atomic_numbers: Sequence[int], model_name: str) -> list[int]: + """Name-keyed species conversion. Prefer ``adapter.to_species``. + + RESIDUAL, and deliberately marked as one. Deciding the species convention + from a *string* is the shape this cluster exists to remove: the caller then + holds a name and the padding values separately, and nothing stops the two + from disagreeing. Every batching caller now asks the adapter instead + (:func:`Auto3D.batch_opt.padding.pad_from_mols`, + :mod:`Auto3D.cli.commands.models`). + + Two call sites remain, both in ``Auto3D.ASE.thermo`` + (``Calculator.calculate`` and ``mol2aimnet_input``), because migrating them + means changing ``Calculator``'s public signature -- work that belongs with + the wider ``thermo.py`` restructuring, not here. Neither builds a padded + batch, so neither can hit the disagreement class; they are a redundant entry + point, not a live defect. Delete this function when they move. + + Args: + atomic_numbers: Atomic numbers, one per atom. + model_name: Engine name, matched case-insensitively against ``"ANI2xt"`` + (the same normalization ``ModelFactory.create`` applies). Only a + match remaps; every other value (AIMNET, any aimnet registry name, + ANI2x, a custom model path) passes through unchanged. + + Returns: + Species values in the model's own convention. + + Raises: + ValueError: ``model_name`` is ANI2xt and an atomic number is out of set. + """ + if model_name.upper() != MODEL_ANI2XT.upper(): + return list(atomic_numbers) + return to_ani2xt_species(atomic_numbers) diff --git a/src/Auto3D/utils/validation.py b/src/Auto3D/utils/validation.py index 4218bf17..0fd7b6d8 100644 --- a/src/Auto3D/utils/validation.py +++ b/src/Auto3D/utils/validation.py @@ -27,7 +27,7 @@ from Auto3D.utils.stereochemistry import count_unspecified_stereo if TYPE_CHECKING: - pass + from Auto3D.config import Auto3DOptions logger = get_logger(__name__) @@ -291,18 +291,22 @@ def check_input(args: Any) -> None: - use_gpu: Whether to use GPU acceleration - isomer_engine: Engine for isomer enumeration ('rdkit' or 'omega') - optimizing_engine: Engine for geometry optimization ('ANI2x', 'ANI2xt', 'AIMNET', or path) - - opt_steps: Number of optimization steps - input_format: Input file format ('smi' or 'sdf') - path: Path to input file - enumerate_isomer: Whether to enumerate stereoisomers + ``opt_steps`` is no longer read here: its minimum lives in + ``Auto3D.config.FIELD_BOUNDS`` and is enforced when the + configuration is constructed, not when it is used. + Returns: None. The function prints recommendations. Raises: GPUError: If GPU is requested but not available. DependencyError: If required dependency not available (OpenEye, TorchANI). - ConfigurationError: If configuration parameters are invalid (opt_steps, engine mismatch). + ConfigurationError: If the optimizing engine cannot represent the input + molecules (charged, or outside the ANI element set). ModelLoadError: If custom NNP cannot be loaded. """ logger.info("Checking input file...") @@ -351,10 +355,12 @@ def check_input(args: Any) -> None: "https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model" ) from e - if int(args.opt_steps) < 10: - raise ConfigurationError( - f"Number of optimization steps cannot be smaller than 10, but received {args.opt_steps}" - ) + # No opt_steps check here any more. It hand-wrote `< 10` while + # Auto3D.config.FIELD_BOUNDS declared ("ge", 1) -- two minimums for one + # option (see that table's comment). FIELD_BOUNDS now declares 10 and is + # enforced by Auto3DOptions.__post_init__ and CLIConfig's _check_bounds, so + # every entry point rejects opt_steps=5 at construction, before this + # function is reached and before any banner prints. # Check the input format if args.input_format == "smi": @@ -527,35 +533,36 @@ def check_sdf_format(args: Any) -> tuple[bool, list[str]]: return ANI, only_aimnet_ids -def check_valid_configuration( - path: str | None = None, - k: int | bool = False, - window: float | bool = False, - use_gpu: bool = True, - gpu_idx: int | list[int] = 0, - optimizing_engine: str = "AIMNET", - isomer_engine: str = "rdkit", - opt_steps: int = 2000, - enumerate_tautomer: bool = False, - tauto_engine: str = "rdkit", -) -> list[str]: - """Validate Auto3D configuration parameters. - - This function checks if the provided configuration parameters are valid - and compatible with each other. +def check_valid_configuration(options: Auto3DOptions) -> list[str]: + """Validate an ``Auto3DOptions`` for things its own construction cannot check. + + Takes the configuration **object**, not a copy of its field names. The + previous signature re-declared ten option names *and their own defaults* + (``optimizing_engine="AIMNET"``, ``isomer_engine="rdkit"``, + ``tauto_engine="rdkit"``, ``opt_steps=2000`` written as a literal rather + than ``DEFAULT_OPT_STEPS``, ...), which made this function a third + configuration schema alongside ``Auto3DOptions`` and ``CLIConfig`` -- one + that could disagree with both about what an unspecified option means, and + that silently never looked at the other eighteen fields. It also forced two + byte-identical ten-keyword marshalling blocks at its only two call sites + (``auto3D.py``'s ``smiles2mols`` and ``workflow.py``'s + ``WorkflowOrchestrator._validate_input``), both of which read every value + straight off an ``Auto3DOptions`` they already had. + + ``Auto3DOptions`` is the authoritative schema, so what belongs here is only + what a dataclass cannot decide for itself: whether the input file exists, + whether a selector was chosen, whether the requested GPU index exists on + *this* machine, whether the engine name resolves in the model registry, and + whether the OpenEye license the chosen engines need is present in the + environment. Everything checkable from the values alone -- numeric bounds + (``FIELD_BOUNDS``, including ``opt_steps >= 10``), the isomer/tautomer + engine whitelists (``ENGINE_CHOICES``), and selector mutual exclusion -- + has already run in ``__post_init__``, so re-checking it here would be the + duplicate this change removes. Args: - path: Path to input file. Must be provided and exist. - k: Number of top conformers to keep. Either k or window must be specified. - window: Energy window in kcal/mol for conformer selection. Either k or window must be specified. - use_gpu: Whether to use GPU acceleration. - gpu_idx: GPU device index or list of indices. - optimizing_engine: Engine for geometry optimization. Must be one of - 'ANI2x', 'ANI2xt', 'AIMNET' or a valid path to a custom model. - isomer_engine: Engine for isomer enumeration. Must be 'rdkit' or 'omega'. - opt_steps: Number of optimization steps. Must be >= 10. - enumerate_tautomer: Whether to enumerate tautomers. - tauto_engine: Engine for tautomer enumeration. Must be 'rdkit' or 'oechem'. + options: The configuration to check. Any object exposing + ``Auto3DOptions``'s attributes works, matching ``check_input``. Returns: List of error messages. Empty list if configuration is valid. @@ -569,6 +576,11 @@ def check_valid_configuration( check_input already raised for this condition (M23). See check_gpu_requested for the full rationale. """ + path = options.path + use_gpu = options.use_gpu + gpu_idx = options.gpu_idx + isomer_engine = options.isomer_engine + errors: list[str] = [] # Check path @@ -578,7 +590,7 @@ def check_valid_configuration( errors.append(f"Input path does not exist: {path}") # Check k and window - if not k and not window: + if not options.k and not options.window: errors.append("Either 'k' or 'window' must be specified for conformer selection.") # Check GPU configuration. Raises immediately rather than appending to @@ -604,29 +616,28 @@ def check_valid_configuration( # is a pure offline dict read against a bundled YAML, so validating costs # nothing. try: - resolve_engine_name(optimizing_engine) + resolve_engine_name(options.optimizing_engine) except ConfigurationError as exc: errors.append(str(exc)) - # Check isomer_engine - valid_isomer_engines = {"rdkit", "omega"} - if isomer_engine.lower() not in valid_isomer_engines: - errors.append(f"isomer_engine must be one of {valid_isomer_engines}. Got: {isomer_engine}") + # No isomer_engine/tauto_engine whitelist here any more: both are in + # Auto3D.config.ENGINE_CHOICES and enforced by Auto3DOptions.__post_init__, + # so an unrecognized value cannot reach this function. The two local + # `valid_*_engines` sets that used to stand here were the third and fourth + # hand-written copies of those whitelists. + # + # The license checks below stay: they are about the *environment*, not about + # the value, so no amount of construction-time validation can answer them. # Check OpenEye license for omega if isomer_engine.lower() == "omega" and "OE_LICENSE" not in os.environ: errors.append("OpenEye license (OE_LICENSE) not found but omega isomer_engine is selected.") - # Check opt_steps - if opt_steps < 10: - errors.append(f"opt_steps must be >= 10. Got: {opt_steps}") - - # Check tautomer configuration - valid_tauto_engines = {"rdkit", "oechem"} - if enumerate_tautomer and tauto_engine.lower() not in valid_tauto_engines: - errors.append(f"tauto_engine must be one of {valid_tauto_engines}. Got: {tauto_engine}") - - if enumerate_tautomer and tauto_engine.lower() == "oechem" and "OE_LICENSE" not in os.environ: + if ( + options.enumerate_tautomer + and options.tauto_engine.lower() == "oechem" + and "OE_LICENSE" not in os.environ + ): errors.append("OpenEye license (OE_LICENSE) not found but oechem tauto_engine is selected.") return errors diff --git a/src/Auto3D/workflow.py b/src/Auto3D/workflow.py index 6af87cbe..7121351f 100644 --- a/src/Auto3D/workflow.py +++ b/src/Auto3D/workflow.py @@ -237,18 +237,7 @@ def _validate_input(self) -> None: # Without this the bad index only surfaces deep inside a spawned worker as # an opaque "no structure converged". check_valid_configuration already # validates the index against torch.cuda.device_count(); reuse it. - config_errors = check_valid_configuration( - path=self.config.path, - k=self.config.k, - window=self.config.window, - use_gpu=self.config.use_gpu, - gpu_idx=self.config.gpu_idx, - optimizing_engine=self.config.optimizing_engine, - isomer_engine=self.config.isomer_engine, - opt_steps=self.config.opt_steps, - enumerate_tautomer=self.config.enumerate_tautomer, - tauto_engine=self.config.tauto_engine, - ) + config_errors = check_valid_configuration(self.config) if config_errors: raise ConfigurationError( "Invalid configuration:\n - " + "\n - ".join(config_errors) diff --git a/src/Auto3D/workflow_workers.py b/src/Auto3D/workflow_workers.py index 3870f9b0..83c86782 100644 --- a/src/Auto3D/workflow_workers.py +++ b/src/Auto3D/workflow_workers.py @@ -24,6 +24,7 @@ from Auto3D.batch_opt.batchopt import optimizing from Auto3D.config import optimizer_worker_indices from Auto3D.isomers import IsomerEngineFactory +from Auto3D.model_factory import create_model from Auto3D.processors import TautomerProcessor from Auto3D.ranking import ranking from Auto3D.utils import create_chunk_meta_names, housekeeping @@ -230,8 +231,20 @@ def progress_cb(event, _q=progress_queue, _job=job): _q.put({**event, "job": _job}) except Exception: pass + # HARD CONSTRAINT: the adapter is built HERE, inside the spawned + # worker, and must stay here. `optimizing` used to construct it + # itself; hoisting construction one frame out (to this function) + # keeps it in the same process, but hoisting it any further -- to + # `workflow.py`, which drives the pool, where these duplicated + # `create_model` calls would look like an obvious cleanup -- pushes + # a device-resident nn.Module, and for AIMNET a live + # AIMNet2Calculator, across the `spawn` boundary. That is either an + # unpicklable-object failure or CUDA re-initialization in the + # parent, and nothing in the signature says so. + adapter = create_model(optimizing_engine, device) optimizer = optimizing(enumerated_sdf, optimized_og, - optimizing_engine, device, opt_config, + adapter=adapter, device=device, + config=opt_config, progress_cb=progress_cb) optimizer.run() diff --git a/tests/helpers_adapter.py b/tests/helpers_adapter.py new file mode 100644 index 00000000..2760eabc --- /dev/null +++ b/tests/helpers_adapter.py @@ -0,0 +1,113 @@ +# tests/helpers_adapter.py +"""One conforming :class:`~Auto3D.models.contract.ModelAdapter` double. + +``EnForce_ANI.__init__`` gates its first argument against the adapter contract +(``Auto3D.models.contract.ModelAdapter``), and ``pad_from_mols`` reads the +species convention *and* both padding sentinels off the same object. Before +this module every test that needed "something adapter-shaped" grew its own +duck-typed class declaring only the members that particular test happened to +exercise, so the gate could not be tightened without six unrelated files going +red -- and the cheapest way to make them green again would have been to weaken +the gate. + +Everything here is a plain Python object: no ``nn.Module``, no weights, no +device traffic, and nothing is loaded or downloaded. +""" +from __future__ import annotations + +from collections.abc import Sequence + +import torch + + +class FakeAdapter: + """A minimal object that satisfies the whole ``ModelAdapter`` contract. + + Args: + coord_pad: Coordinate fill value reported to the padder. + species_pad: Species fill value reported to the padder. + species_map: Optional atomic-number -> model-species mapping applied by + :meth:`to_species`. ``None`` means the identity, which is what every + adapter except ``ANI2xtAdapter`` does. + energy_fn: Optional ``(coords, species, charges) -> energies``. The + default is ``sum(coords**2)`` per molecule, whose gradient is + analytic (``2*coords``), so a caller can check forces without any + model. + """ + + def __init__( + self, + coord_pad: float = 0.0, + species_pad: int = -1, + species_map: dict[int, int] | None = None, + energy_fn=None, + ) -> None: + self.coord_pad = coord_pad + self.species_pad = species_pad + self.species_map = species_map + self._energy_fn = energy_fn + #: Recorded ``(coords_dtype, species, charges)`` per forward/energy call. + self.calls: list[dict] = [] + + # -- the species half of the contract --------------------------------- + def to_species(self, atomic_numbers: Sequence[int]) -> list[int]: + if self.species_map is None: + return list(atomic_numbers) + return [self.species_map[int(z)] for z in atomic_numbers] + + # -- the numerical half ----------------------------------------------- + def _energies(self, coords: torch.Tensor, species, charges) -> torch.Tensor: + if self._energy_fn is not None: + return self._energy_fn(coords, species, charges) + return coords.pow(2).sum(dim=(1, 2)) + + def forward(self, coords, species, charges, atom_mask=None): + self.calls.append( + {"dtype": coords.dtype, "atom_mask": atom_mask, "kind": "forward"} + ) + coords = coords if coords.requires_grad else coords.detach().requires_grad_(True) + energy = self._energies(coords, species, charges) + grad = torch.autograd.grad([energy.sum()], [coords], create_graph=False)[0] + return energy, -grad + + def energy(self, coords, species, charges, atom_mask=None): + """Energies at the dtype of ``coords`` -- deliberately no downcast.""" + self.calls.append( + {"dtype": coords.dtype, "atom_mask": atom_mask, "kind": "energy"} + ) + return self._energies(coords, species, charges) + + +class AdapterModuleMixin: + """Makes an ``nn.Module`` test double satisfy ``ModelAdapter``. + + For doubles that must be real ``nn.Module``s (because the code under test + reads ``.parameters()``, or ASE needs a module) and therefore cannot simply + be :class:`FakeAdapter`. Mix in FIRST so these defaults are found before + ``nn.Module``'s attribute machinery:: + + class _Stub(AdapterModuleMixin, nn.Module): + def forward(self, coords, species, charges, atom_mask=None): ... + + The values match ``BaseModelAdapter``'s own defaults. ``species_pad = -1`` + specifically: it can be neither a real atomic number nor a 0-based species + index, so it cannot collide the way ``0`` did (audit C13). + """ + + coord_pad: float = 0.0 + species_pad: int = -1 + + def to_species(self, atomic_numbers: Sequence[int]) -> list[int]: + return list(atomic_numbers) + + def energy(self, coords, species, charges, atom_mask=None): + return self.forward(coords, species, charges, atom_mask)[0] + + +def padded_batch(n_mols: int = 2, n_atoms: int = 3): + """Tensors shaped like :func:`Auto3D.batch_opt.padding.pad_from_mols`.""" + coords = torch.zeros(n_mols, n_atoms, 3) + species = torch.ones(n_mols, n_atoms, dtype=torch.long) + charges = torch.zeros(n_mols) + atom_mask = torch.ones(n_mols, n_atoms, dtype=torch.bool) + return coords, species, charges, atom_mask diff --git a/tests/test_SPE.py b/tests/test_SPE.py index ddee4b83..e08373a5 100644 --- a/tests/test_SPE.py +++ b/tests/test_SPE.py @@ -297,8 +297,10 @@ def forward_batched(self, coords, numbers, charges, atom_mask=None): pad_calls = [] - def fake_pad(mols, model_name, device, coord_pad, species_pad): - pad_calls.append((coord_pad, species_pad)) + def fake_pad(mols, adapter, device): + # The padder now reads BOTH sentinels off the adapter it was handed, so + # they cannot come from two places and disagree (audit C3/C4). + pad_calls.append((adapter.coord_pad, adapter.species_pad)) n = len(mols) return ( torch.zeros(n, 1, 3), diff --git a/tests/test_adapter_atom_mask.py b/tests/test_adapter_atom_mask.py index 946e6cfe..7cf50151 100644 --- a/tests/test_adapter_atom_mask.py +++ b/tests/test_adapter_atom_mask.py @@ -23,6 +23,14 @@ from torch import nn from Auto3D.batch_opt.padding import pad_from_mols +from tests.helpers_adapter import FakeAdapter + +# Stands in for AIMNet2's padding convention (raw atomic numbers, species_pad=0) +# without loading anything. `pad_from_mols` reads the species convention AND both +# fill values off this one object, so they cannot disagree (audit C3/C4). +def _aimnet_padding(): + return FakeAdapter(coord_pad=0.0, species_pad=0) + from Auto3D.models.adapter import AIMNet2Adapter rdkit = pytest.importorskip("rdkit") @@ -90,9 +98,7 @@ class TestDummyAtomIsNotTreatedAsPadding: @staticmethod def _pad(mols): - return pad_from_mols( - mols, "AIMNET", torch.device("cpu"), coord_pad=0.0, species_pad=0 - ) + return pad_from_mols(mols, _aimnet_padding(), torch.device("cpu")) def test_r_group_atom_is_counted(self): mol = _embed("*CCO", "rgroup") @@ -135,8 +141,7 @@ def test_padded_slots_are_dropped_and_get_zero_force(self): big = _embed("*CCO", "rgroup") # 9 atoms small = _embed("O", "water") # 3 atoms coords, species, charges, atom_mask = pad_from_mols( - [big, small], "AIMNET", torch.device("cpu"), - coord_pad=0.0, species_pad=0, + [big, small], _aimnet_padding(), torch.device("cpu") ) assert species.shape == (2, 9) assert int(atom_mask.sum()) == 12 @@ -159,7 +164,7 @@ class TestUnpaddedCallersNeedNoMask: def test_no_mask_treats_every_slot_as_real(self): mol = _embed("*CCO", "rgroup") coords, species, charges, _ = pad_from_mols( - [mol], "AIMNET", torch.device("cpu"), coord_pad=0.0, species_pad=0 + [mol], _aimnet_padding(), torch.device("cpu") ) calc = _RecordingCalculator() energy, _ = _stub_adapter(calc).forward(coords, species, charges) @@ -175,8 +180,9 @@ def test_forward_batched_forwards_the_mask_per_sub_batch(self): seen: list[int] = [] - class _CountingAdapter(nn.Module): - coord_pad = 0.0 + from tests.helpers_adapter import AdapterModuleMixin + + class _CountingAdapter(AdapterModuleMixin, nn.Module): species_pad = 0 def forward(self, coords, species, charges, atom_mask=None): @@ -189,8 +195,7 @@ def forward(self, coords, species, charges, atom_mask=None): big = _embed("*CCO", "rgroup") small = _embed("O", "water") coords, species, charges, atom_mask = pad_from_mols( - [big, small], "AIMNET", torch.device("cpu"), - coord_pad=0.0, species_pad=0, + [big, small], _aimnet_padding(), torch.device("cpu") ) # batchsize_atoms=9 with N=9 gives one molecule per sub-batch, so the # mask has to be sliced with the same indices as coord/numbers. diff --git a/tests/test_batchopt.py b/tests/test_batchopt.py index bbabd78c..0143f917 100644 --- a/tests/test_batchopt.py +++ b/tests/test_batchopt.py @@ -2,60 +2,13 @@ """Unit tests for the batchopt module.""" from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest import torch from Auto3D.batch_opt.batchopt import optimizing, EnForce_ANI - - -class TestOptimizingUsesModelFactory: - """Tests for optimizing class using ModelFactory.""" - - def test_optimizing_uses_model_factory(self): - """optimizing class should use ModelFactory for model creation.""" - with patch('Auto3D.batch_opt.batchopt.create_model') as mock_factory: - mock_adapter = MagicMock() - mock_adapter.coord_pad = 0.0 - mock_adapter.species_pad = 0 - mock_factory.return_value = mock_adapter - - config = { - 'opt_steps': 100, - 'opttol': 0.003, - 'patience': 1000, - 'batchsize_atoms': 1024 - } - device = torch.device("cpu") - opt = optimizing("dummy.sdf", "out.sdf", "AIMNET", device, config) - - # Check that create_model was called with the right model name and device - mock_factory.assert_called_once_with("AIMNET", device) - # Verify the adapter's properties are used - assert opt.coord_pad == 0.0 - assert opt.species_pad == 0 - - def test_optimizing_uses_adapter_padding_values(self): - """optimizing should get coord_pad and species_pad from the adapter.""" - with patch('Auto3D.batch_opt.batchopt.create_model') as mock_factory: - mock_adapter = MagicMock() - mock_adapter.coord_pad = 1.5 - mock_adapter.species_pad = -2 - mock_factory.return_value = mock_adapter - - config = { - 'opt_steps': 100, - 'opttol': 0.003, - 'patience': 1000, - 'batchsize_atoms': 1024 - } - device = torch.device("cpu") - opt = optimizing("dummy.sdf", "out.sdf", "AIMNET", device, config) - - # Verify padding values come from adapter - assert opt.coord_pad == 1.5 - assert opt.species_pad == -2 +from tests.helpers_adapter import FakeAdapter class TestEnForceANI: @@ -182,10 +135,12 @@ def test_oscillating_structure_is_not_reported_converged(self, job_dir): # patience=1 guarantees the oscillation path is taken on any structure # that does not reduce fmax on its very first step. + from Auto3D.model_factory import create_model + opt = optimizing( in_f=str(sdf), out_f=str(job_dir / "out.sdf"), - name="ANI2xt", + adapter=create_model("ANI2xt", torch.device("cpu")), device=torch.device("cpu"), config={"opt_steps": 5, "opttol": 1e-9, "patience": 1, "batchsize_atoms": 1024}, ) @@ -219,20 +174,11 @@ def test_oscillating_structure_is_not_reported_converged(self, job_dir): def test_make_buckets_groups_by_size(tmp_path, monkeypatch): """Buckets must be size-homogeneous; a size outlier splits into its own bucket.""" - from types import SimpleNamespace - from rdkit import Chem from rdkit.Chem import AllChem import torch from Auto3D.batch_opt.batchopt import optimizing - # _make_buckets is pure-Python; stub create_model so this never loads the - # real AIMNet2 model. - monkeypatch.setattr( - "Auto3D.batch_opt.batchopt.create_model", - lambda *a, **k: SimpleNamespace(coord_pad=0.0, species_pad=-1), - ) - # Build an optimizing instance without running (just to call _make_buckets) inp = tmp_path / "in.sdf" sizes = ["C", "CC", "CCC", "C1CCCCCCCCCCCCCCCCCCC1"] # tiny ... and one big ring @@ -241,8 +187,12 @@ def test_make_buckets_groups_by_size(tmp_path, monkeypatch): for i, s in enumerate(sizes): m = Chem.AddHs(Chem.MolFromSmiles(s)); AllChem.EmbedMolecule(m, randomSeed=1) m.SetProp("_Name", str(i)); w.write(m); mols.append(m) - eng = optimizing(str(inp), str(tmp_path/"o.sdf"), "AIMNET", torch.device("cpu"), - {"opt_steps":1,"opttol":0.01,"patience":1,"batchsize_atoms":1024}) + # _make_buckets is pure-Python, so a conforming double is enough and no + # model is loaded. `optimizing` no longer builds its own adapter, so there is + # no create_model seam left to patch. + eng = optimizing(str(inp), str(tmp_path/"o.sdf"), adapter=FakeAdapter(), + device=torch.device("cpu"), + config={"opt_steps":1,"opttol":0.01,"patience":1,"batchsize_atoms":1024}) buckets = eng._make_buckets(mols) # the big 20-carbon ring must not share a bucket with methane big_idx = 3 @@ -252,8 +202,6 @@ def test_make_buckets_groups_by_size(tmp_path, monkeypatch): def test_optimizing_preserves_input_order(tmp_path, monkeypatch): """Bucketing reorders internally but output order must match input.""" - from types import SimpleNamespace - from rdkit import Chem from rdkit.Chem import AllChem import torch @@ -274,16 +222,99 @@ def fake_ensemble_opt(net, coord, numbers, charges, param, device, numbers=numbers.tolist(), converged_mask=[True]*n, oscillating_count=[0]*n) monkeypatch.setattr(bo, "ensemble_opt", fake_ensemble_opt) - # The optimization itself is faked above; stub create_model so constructing - # `optimizing` does not load the real AIMNet2 model. - monkeypatch.setattr( - bo, "create_model", - lambda *a, **k: SimpleNamespace(coord_pad=0.0, species_pad=-1), - ) out = tmp_path / "out.sdf" - eng = bo.optimizing(str(inp), str(out), "AIMNET", torch.device("cpu"), - {"opt_steps":1,"opttol":0.01,"patience":1,"batchsize_atoms":1024}) + eng = bo.optimizing(str(inp), str(out), adapter=FakeAdapter(), + device=torch.device("cpu"), + config={"opt_steps":1,"opttol":0.01,"patience":1,"batchsize_atoms":1024}) eng.run() names = [m.GetProp("_Name") for m in Chem.SDMolSupplier(str(out), removeHs=False)] assert names == ["0", "1", "2"] # original input order + + +class TestBatchOptDependsDownwards: + """``batch_opt`` must depend on ``models/``, never on ``model_factory``. + + ``batchopt.py`` imported ``Auto3D.model_factory.create_model`` at module + scope and called it in ``optimizing.__init__``: the numerical layer + constructing its own dependency, and reaching UP into the layer that is + supposed to sit above it. The visible symptom was in the tests -- every one + of them had to monkeypatch ``Auto3D.batch_opt.batchopt.create_model``, a seam + that existed only because the arrow pointed the wrong way. + """ + + def test_importing_batchopt_does_not_pull_in_the_factory(self): + """Asserted in a fresh interpreter: an already-imported ``model_factory`` + would make this vacuous inside the test session.""" + import subprocess + import sys + + program = ( + "import sys; import Auto3D.batch_opt.batchopt as b; " + "assert 'Auto3D.model_factory' not in sys.modules, " + "sorted(m for m in sys.modules if m.startswith('Auto3D')); " + "print('ok')" + ) + result = subprocess.run( + [sys.executable, "-c", program], capture_output=True, text=True + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "ok" in result.stdout + + def test_batchopt_exposes_no_create_model_seam(self): + import Auto3D.batch_opt.batchopt as bo + + assert not hasattr(bo, "create_model") + + +class TestOptimizingTakesAnAdapterNotAName: + """Construction is the caller's job, and the caller must be in the worker. + + ``optimizing`` no longer knows an engine name at all: ``pad_from_mols`` asks + the adapter, so there is nothing left for a name to decide. + """ + + @staticmethod + def _config(): + return { + "opt_steps": 100, + "opttol": 0.003, + "patience": 1000, + "batchsize_atoms": 1024, + } + + def test_an_engine_name_is_rejected(self): + """The parameters after ``out_f`` are keyword-only, so a stale positional + call fails at the call rather than silently binding a string into the + slot that supplies the padding values.""" + with pytest.raises(TypeError): + optimizing( + "dummy.sdf", "out.sdf", "AIMNET", torch.device("cpu"), self._config() + ) + + def test_padding_values_come_from_the_injected_adapter(self): + from tests.helpers_adapter import FakeAdapter + + adapter = FakeAdapter(coord_pad=1.5, species_pad=-2) + opt = optimizing( + "dummy.sdf", + "out.sdf", + adapter=adapter, + device=torch.device("cpu"), + config=self._config(), + ) + assert opt.model is adapter + assert opt.coord_pad == 1.5 + assert opt.species_pad == -2 + + def test_optimizing_no_longer_carries_an_engine_name(self): + from tests.helpers_adapter import FakeAdapter + + opt = optimizing( + "dummy.sdf", + "out.sdf", + adapter=FakeAdapter(), + device=torch.device("cpu"), + config=self._config(), + ) + assert not hasattr(opt, "name") diff --git a/tests/test_cli_config_schema.py b/tests/test_cli_config_schema.py index 047c5dfb..ae1e3b84 100644 --- a/tests/test_cli_config_schema.py +++ b/tests/test_cli_config_schema.py @@ -269,32 +269,66 @@ def test_shipped_parameters_yaml_loads(): cfg.to_auto3d_options() # must not raise -def test_shipped_legacy_v2_parameters_yaml_loads(): - """docs/legacy-v2/parameters.yaml (``k: 1`` / ``window: False``) must - validate through the exact construction ``auto3Dcli._run_legacy_yaml`` - uses -- ``yaml.safe_load`` + the "None"-string-to-None conversion + - ``CLIConfig(**parameters)`` (auto3Dcli.py, around the ``CLIConfig( - **parameters)`` call) -- not the pipeline itself. Before this fix, - ``window: False`` was coerced by Pydantic to ``0.0`` ahead of - ``CLIConfig``'s bound-check model validator, which then rejected it as - a non-positive window: this exact file, run through this exact CLI - entry point, raised ``ValidationError`` and exited 1 on this branch - while working unmodified on `main`. +def test_shipped_parameters_yaml_is_complete(): + """The shipped example must show every option, or it teaches a subset. + + ``parameters.yaml`` is an *instance*, not a fourth schema -- but an instance + is how most users discover which options exist, and three + (``use_parallel_embedding``, ``parallel_workers``, + ``parallel_embedding_threshold``) were simply absent with nothing noticing. + Anything deliberately omitted goes on the allowlist below *with a reason*, + so "missing" and "intentionally not shown" stop being the same state. """ import yaml as yaml_mod - from Auto3D.cli.config_schema import CLIConfig + from Auto3D.cli.config_schema import OPTIONS_ONLY_FIELDS + + omitted = { + # k is set instead; check_selectors_mutually_exclusive rejects both at + # once, so an example cannot demonstrate the two together. The key is + # still present as `window: None` for discoverability. + "window": "mutually exclusive with k, which the example sets", + } repo_root = Path(__file__).resolve().parent.parent - yaml_path = repo_root / "docs" / "legacy-v2" / "parameters.yaml" + data = yaml_mod.safe_load((repo_root / "parameters.yaml").read_text()) + + expected = set(_auto3d_option_fields()) - set(OPTIONS_ONLY_FIELDS) + missing = expected - set(data) - set(omitted) + assert not missing, ( + f"parameters.yaml does not mention these options: {sorted(missing)}. " + f"Add them, or add them to this test's `omitted` allowlist with a reason." + ) + unknown = set(data) - expected + assert not unknown, f"parameters.yaml sets options that do not exist: {unknown}" + + +def test_shipped_legacy_v2_parameters_yaml_loads(): + """docs/legacy-v2/parameters.yaml (``k: 1`` / ``window: False``) must + validate through ``load_yaml_config`` -- the function ``auto3Dcli. + _run_legacy_yaml`` and ``cli/commands/run.py`` both call, and now the + only YAML ingestion path there is -- not through the pipeline itself. + Before this fix, ``window: False`` was coerced by Pydantic to ``0.0`` + ahead of ``CLIConfig``'s bound-check model validator, which then rejected + it as a non-positive window: this exact file, run through this exact CLI + entry point, raised ``ValidationError`` and exited 1 on this branch while + working unmodified on `main`. + + This used to re-implement the legacy path inline (``yaml.safe_load`` + the + "None"-string-to-None loop + ``CLIConfig(**parameters)``) while claiming to + use "the exact construction ``_run_legacy_yaml`` uses". That claim was what + made the duplicate ingestion layer invisible: the copy under test could + only ever agree with the copy in ``auto3Dcli.py``, so the three shape + guards the legacy path was missing (empty file, non-mapping top level, YAML + syntax error) were untested from either side. Calling the real function is + strictly better -- it tests the path instead of a replica of it. + """ + from Auto3D.cli.config_schema import load_yaml_config - with open(yaml_path) as f: - parameters = yaml_mod.safe_load(f) - for key, val in list(parameters.items()): - if val == "None": - parameters[key] = None + repo_root = Path(__file__).resolve().parent.parent + yaml_path = repo_root / "docs" / "legacy-v2" / "parameters.yaml" - config = CLIConfig(**parameters) # must not raise + config = load_yaml_config(yaml_path) # must not raise assert config.k == 1 assert config.window is None # False normalized to CLIConfig's own sentinel assert config.memory is None @@ -369,17 +403,312 @@ def test_build_cli_config_translates_non_pydantic_validator_errors(): assert build_cli_config(path=Path("in.smi"), k=1, gpu_idx="0,1").gpu_idx == [0, 1] -def test_cliconfig_covers_all_auto3doptions_fields(): - """Guard against config-layer drift: every user-facing Auto3DOptions field - must be reachable from the CLI/YAML via CLIConfig. ``input_format`` is set - internally by the workflow, so it is the only allowed exclusion.""" +# ============================================================================= +# CLIConfig <-> Auto3DOptions parity +# +# `Auto3DOptions` (Auto3D/config.py) is the authoritative configuration schema: +# the Python API's `main()`/`smiles2mols` take it, so it cannot depend on the +# CLI layer, which makes it the only candidate for "source". `CLIConfig` stays +# hand-written -- generating it with `pydantic.create_model` would hide the +# fields from mypy and IDEs and couple the config layer to a metaprogramming API +# -- and the four tests below make drift impossible to *merge* instead of +# impossible to *write*. Adding an option is two edits (dataclass field, +# CLIConfig field); the second is named by a failing test until it is made. +# +# Four legs, because each catches a failure the others cannot: +# 1. names, both directions +# 2. defaults +# 3. sentinels (None here, False there) +# 4. round-trip through to_auto3d_options() <-- nothing checked this before +# ============================================================================= + + +def _auto3d_option_fields(): import dataclasses - from Auto3D.cli.config_schema import CLIConfig from Auto3D.config import Auto3DOptions - excluded = {"input_format"} - opt_fields = {f.name for f in dataclasses.fields(Auto3DOptions)} - excluded + return {f.name: f for f in dataclasses.fields(Auto3DOptions)} + + +def test_cliconfig_covers_all_auto3doptions_fields(): + """Leg 1: name parity, in BOTH directions. + + Every user-facing ``Auto3DOptions`` field must be reachable from the + CLI/YAML via ``CLIConfig``, and ``CLIConfig`` must not carry a field the + core schema has never heard of -- an option a user can set in a YAML file + and that ``to_auto3d_options`` then has nowhere to put. Only the reachability + direction was checked before. + + The exclusion list is ``config_schema.OPTIONS_ONLY_FIELDS``, which states + the reason per field, rather than a set literal repeated here. + """ + from Auto3D.cli.config_schema import OPTIONS_ONLY_FIELDS, CLIConfig + + opt_fields = set(_auto3d_option_fields()) - set(OPTIONS_ONLY_FIELDS) cli_fields = set(CLIConfig.model_fields) - missing = opt_fields - cli_fields - assert not missing, f"CLIConfig is missing Auto3DOptions fields: {missing}" + + assert not opt_fields - cli_fields, ( + f"CLIConfig is missing Auto3DOptions fields: {opt_fields - cli_fields}" + ) + assert not cli_fields - opt_fields, ( + f"CLIConfig declares fields Auto3DOptions does not have: " + f"{cli_fields - opt_fields}. A user could set them and nothing would " + f"read them." + ) + # An excluded field must actually exist on Auto3DOptions, or the exclusion + # list is silently hiding a typo instead of an internal field. + assert set(OPTIONS_ONLY_FIELDS) <= set(_auto3d_option_fields()) + + +def test_cliconfig_defaults_match_auto3doptions(): + """Leg 2: default parity. + + A default written twice is a default that can differ by entry point -- + ``auto3d run in.smi`` and ``main(Auto3DOptions(path=...))`` would then be + different runs. Sentinel fields are excluded here and checked by leg 3, + which is about the deliberate ``None``/``False`` difference. + """ + from Auto3D.cli.config_schema import OPTIONS_ONLY_FIELDS, CLIConfig + from Auto3D.config import SENTINEL_FIELDS + + skip = set(OPTIONS_ONLY_FIELDS) | set(SENTINEL_FIELDS) + mismatched = {} + for name, spec in _auto3d_option_fields().items(): + if name in skip: + continue + cli_default = CLIConfig.model_fields[name].default + if cli_default != spec.default: + mismatched[name] = (spec.default, cli_default) + assert not mismatched, ( + f"default drift between Auto3DOptions and CLIConfig " + f"(field: (Auto3DOptions, CLIConfig)): {mismatched}" + ) + + +def test_sentinel_fields_use_the_expected_sentinel_on_each_side(): + """Leg 3: sentinel parity. + + The two classes spell "not specified" differently on purpose -- ``None`` on + ``CLIConfig`` (pydantic coerces ``False`` to ``0`` before any bound check; + see ``_false_means_unset``) and ``False``/``None`` on ``Auto3DOptions``. That + difference is bridged, and a bridge is only safe while both ends are known, + so both ends are asserted rather than assumed. + """ + from Auto3D.cli.config_schema import CLIConfig + from Auto3D.config import SENTINEL_FIELDS + + for name in sorted(SENTINEL_FIELDS): + assert CLIConfig.model_fields[name].default is None, name + assert _auto3d_option_fields()[name].default in (False, None), name + + +# One non-default value per CLIConfig field, for the round-trip leg. Every value +# differs from the field's default, so a field that fails to cross shows up as +# the default rather than as an equal value. `optimizing_engine` is 'ANI2xt', a +# built-in name that short-circuits registry resolution without importing the +# optional `aimnet` package (the same reason tests/test_cli.py picks it). +_ROUND_TRIP_VALUES: dict[str, object] = { + "path": Path("elsewhere/other.smi"), + "verbose": True, + "job_name": "round-trip", + "enumerate_tautomer": True, + "tauto_engine": "oechem", + "pKaNorm": False, + "isomer_engine": "omega", + "enumerate_isomer": False, + "mode_oe": "dense", + "max_confs": 11, + "mpi_np": 3, + "optimizing_engine": "ANI2xt", + "use_gpu": False, + "gpu_idx": 2, + "opt_steps": 777, + "convergence_threshold": 0.02, + "patience": 111, + "threshold": 0.44, + "batchsize_atoms": 2048, + "use_parallel_embedding": True, + "parallel_workers": 5, + "parallel_embedding_threshold": 11, + "memory": 12, + "capacity": 43, + "allow_tf32": True, +} + +# k and window are mutually exclusive, so they cannot both be set in one object; +# the round-trip runs once per selector instead. +_ROUND_TRIP_SELECTORS: dict[str, object] = {"k": 7, "window": 4.5} + + +def test_round_trip_values_table_covers_every_field(): + """The round-trip table must name every field, or leg 4 silently shrinks. + + A new option whose value is missing here would be round-tripped at its + default, which the assertion below could never distinguish from "did not + arrive". This is the test that fails first when an option is added. + """ + from Auto3D.cli.config_schema import CLIConfig + + covered = set(_ROUND_TRIP_VALUES) | set(_ROUND_TRIP_SELECTORS) + assert covered == set(CLIConfig.model_fields), ( + f"missing from the round-trip table: {set(CLIConfig.model_fields) - covered}; " + f"unknown fields in it: {covered - set(CLIConfig.model_fields)}" + ) + # Every listed value must really differ from the default, or the assertion + # in the round-trip test degenerates into "the default equals the default". + same_as_default = { + name: value + for name, value in _ROUND_TRIP_VALUES.items() + if value == CLIConfig.model_fields[name].default + } + assert not same_as_default, ( + f"these round-trip values equal the field default, so they cannot " + f"detect a dropped field: {same_as_default}" + ) + + +@pytest.mark.parametrize("selector", sorted(_ROUND_TRIP_SELECTORS)) +def test_to_auto3d_options_forwards_every_field(selector): + """Leg 4: round-trip parity -- the leg nothing checked before this. + + ``to_auto3d_options`` used to be 27 hand-written ``field=self.field`` + assignments, and the only parity test compared field-*name* sets. Deleting + one assignment therefore passed every test while silently dropping a user's + setting: ``auto3d run in.smi -c cfg.yaml`` would read ``patience: 50`` from + the file, validate it, print it, and then run with 250. + + It is now a ``dataclasses.fields(Auto3DOptions)`` loop with an explicit + transform table for the four fields whose value genuinely differs across the + boundary, so forwarding is structural. This test is what keeps it structural. + """ + from Auto3D.cli.config_schema import CLIConfig + + other = next(s for s in _ROUND_TRIP_SELECTORS if s != selector) + config = CLIConfig( + **_ROUND_TRIP_VALUES, **{selector: _ROUND_TRIP_SELECTORS[selector]} + ) + options = config.to_auto3d_options() + + # The chosen selector crosses; the other stays "not specified", spelled + # False on the Auto3DOptions side. + assert options[selector] == _ROUND_TRIP_SELECTORS[selector] + assert options[other] is False + + dropped = {} + for name, value in _ROUND_TRIP_VALUES.items(): + arrived = options[name] + # `path` is the one type change: Path -> str (str(None) would be the + # literal "None", so the transform is not just `str`). + expected = str(value) if name == "path" else value + if arrived != expected: + dropped[name] = (expected, arrived) + assert not dropped, ( + f"to_auto3d_options did not forward these fields " + f"(field: (sent, arrived)): {dropped}" + ) + + +def test_engine_choices_table_matches_cliconfig_literals(): + """``Auto3D.config.ENGINE_CHOICES`` is the isomer/tautomer whitelist. + + ``CLIConfig`` keeps ``Literal`` annotations because a ``Literal`` is a type + -- mypy and pydantic's error messages both use it -- but it is the one + remaining hand-written copy of that information, so it is pinned here. This + is deliberately NOT the ``Field(ge=)`` duplication ``FIELD_BOUNDS``'s + docstring forbids: a ``Literal`` cannot be derived from the table without + losing static typing, whereas a numeric bound in a ``Field`` buys nothing a + type can use. + """ + import typing + + from Auto3D.cli.config_schema import CLIConfig + from Auto3D.config import ENGINE_CHOICES + + for name, choices in ENGINE_CHOICES.items(): + annotation = CLIConfig.model_fields[name].annotation + assert typing.get_args(annotation) == choices, ( + f"{name}: CLIConfig Literal{typing.get_args(annotation)} disagrees " + f"with Auto3D.config.ENGINE_CHOICES{choices}" + ) + + +def test_config_init_tables_only_name_real_options(): + """``auto3d config init``'s three tables are instances, and must stay so. + + ``cli/commands/config.py`` holds ``DEFAULT_CONFIG`` (the template), ``PRESETS`` + (quick/balanced/thorough) and ``generate_commented_yaml``'s ``comments`` -- + hand-written key lists, like ``parameters.yaml``. They are not schemas and are + not consolidated away: a template is a curated subset by design, and each + comment is CLI-facing prose, not the field documentation ``Auto3DOptions`` + already carries in its per-field docstrings. What they must never do is drift + into naming an option that does not exist (silently ignored, or rejected by + ``extra="forbid"`` the moment a user runs the file they were just given) or + emit a template key with no explanation. + """ + from Auto3D.cli.commands.config import DEFAULT_CONFIG, PRESETS, generate_commented_yaml + from Auto3D.cli.config_schema import OPTIONS_ONLY_FIELDS + + real = set(_auto3d_option_fields()) - set(OPTIONS_ONLY_FIELDS) + + assert not set(DEFAULT_CONFIG) - real, ( + f"DEFAULT_CONFIG names options that do not exist: " + f"{set(DEFAULT_CONFIG) - real}" + ) + for name, preset in PRESETS.items(): + assert not set(preset) - real, f"preset {name!r}: {set(preset) - real}" + + # Every template key must come out commented, so `config init` never hands a + # user a bare key they have to look up elsewhere. + lines = generate_commented_yaml(dict(DEFAULT_CONFIG)).splitlines() + uncommented = [] + for key in DEFAULT_CONFIG: + index = next( + (i for i, line in enumerate(lines) if line.startswith(f"{key}:")), None + ) + if index is None or index == 0 or not lines[index - 1].startswith("#"): + uncommented.append(key) + assert not uncommented, ( + f"`auto3d config init` emits these keys with no explanatory comment: " + f"{uncommented}" + ) + + # And the template must itself be a runnable configuration. + from Auto3D.cli.config_schema import build_cli_config + + build_cli_config(**DEFAULT_CONFIG) + + +def test_no_field_bounds_field_declares_a_second_pydantic_constraint(): + """``FIELD_BOUNDS`` is the only place a numeric bound may be declared. + + Its docstring says so, ``_check_bounds`` enforces it for every field on both + entry points, and a previous change that added ``Field(ge=1)`` to fields + already in the table had to be reverted -- two constraint sets for one + option is precisely the defect this module's parity tests exist to prevent, + and a pydantic constraint also fails with a different exception path than + ``check_field_bounds``'s ``ConfigurationError``. + + Checked against pydantic's own metadata rather than the source text, so it + holds however the constraint is spelled (``Field(ge=...)``, + ``Annotated[int, Ge(...)]``, ``conint``, ...). + """ + from Auto3D.cli.config_schema import CLIConfig + from Auto3D.config import FIELD_BOUNDS + + forbidden = ("ge", "gt", "le", "lt", "multiple_of", "allow_inf_nan") + offenders = {} + for name in FIELD_BOUNDS: + constraints = [ + f"{attr}={getattr(meta, attr)}" + for meta in CLIConfig.model_fields[name].metadata + for attr in forbidden + if getattr(meta, attr, None) is not None + ] + if constraints: + offenders[name] = constraints + assert not offenders, ( + f"these CLIConfig fields declare a numeric constraint that already " + f"lives in Auto3D.config.FIELD_BOUNDS: {offenders}. Bounds go in that " + f"table only; _check_bounds applies them to every entry point." + ) diff --git a/tests/test_cli_exit_codes.py b/tests/test_cli_exit_codes.py index b98e5167..4dad6681 100644 --- a/tests/test_cli_exit_codes.py +++ b/tests/test_cli_exit_codes.py @@ -469,7 +469,12 @@ def test_exit_5_unloadable_custom_model(tmp_path): def test_exit_5_non_finite_model_output(monkeypatch): """The other ``ModelError`` subclass: a model that loads but produces NaN.""" - class _NanAdapter: + # AdapterModuleMixin supplies the ModelAdapter members this stub does not + # care about; `models test` asks the adapter for the species convention now + # instead of resolving it from the engine name separately (audit C4). + from tests.helpers_adapter import AdapterModuleMixin + + class _NanAdapter(AdapterModuleMixin): def forward(self, coords, species, charges): return torch.tensor([float("nan")]), torch.zeros(1, 5, 3) diff --git a/tests/test_cli_property_commands.py b/tests/test_cli_property_commands.py index 6bb26c6f..66a58d22 100644 --- a/tests/test_cli_property_commands.py +++ b/tests/test_cli_property_commands.py @@ -298,7 +298,12 @@ def test_models_test_success(monkeypatch): """`models test` loads the engine and runs a forward; reports success.""" import torch - class _StubAdapter: + # AdapterModuleMixin supplies the ModelAdapter members this stub does not + # care about; `models test` asks the adapter for the species convention now + # instead of resolving it from the engine name separately (audit C4). + from tests.helpers_adapter import AdapterModuleMixin + + class _StubAdapter(AdapterModuleMixin): def forward(self, coords, species, charges): return torch.zeros(1), torch.zeros(1, 5, 3) @@ -327,7 +332,12 @@ def test_models_test_non_finite_exit_code(monkeypatch): """Non-finite outputs are reported as a model (numerical) error -> exit 5.""" import torch - class _NanAdapter: + # AdapterModuleMixin supplies the ModelAdapter members this stub does not + # care about; `models test` asks the adapter for the species convention now + # instead of resolving it from the engine name separately (audit C4). + from tests.helpers_adapter import AdapterModuleMixin + + class _NanAdapter(AdapterModuleMixin): def forward(self, coords, species, charges): return torch.tensor([float("nan")]), torch.zeros(1, 5, 3) @@ -357,7 +367,12 @@ def test_models_test_no_gpu_still_works_without_cuda(monkeypatch): """--no-gpu must keep working on a CPU-only box (not a blanket failure).""" import torch - class _StubAdapter: + # AdapterModuleMixin supplies the ModelAdapter members this stub does not + # care about; `models test` asks the adapter for the species convention now + # instead of resolving it from the engine name separately (audit C4). + from tests.helpers_adapter import AdapterModuleMixin + + class _StubAdapter(AdapterModuleMixin): def forward(self, coords, species, charges): return torch.zeros(1), torch.zeros(1, 5, 3) @@ -372,7 +387,12 @@ def test_models_test_gpu_works_when_cuda_present(monkeypatch): """--gpu (the default) must still succeed when CUDA is actually available.""" import torch - class _StubAdapter: + # AdapterModuleMixin supplies the ModelAdapter members this stub does not + # care about; `models test` asks the adapter for the species convention now + # instead of resolving it from the engine name separately (audit C4). + from tests.helpers_adapter import AdapterModuleMixin + + class _StubAdapter(AdapterModuleMixin): def forward(self, coords, species, charges): return torch.zeros(1), torch.zeros(1, 5, 3) diff --git a/tests/test_config.py b/tests/test_config.py index 27376a07..776d0a7a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -2,6 +2,7 @@ from __future__ import annotations from dataclasses import replace +from pathlib import Path import pytest @@ -185,39 +186,58 @@ def test_capacity_default_matches_across_layers(): assert Auto3DOptions(path="x.smi").capacity == CLIConfig(path="x.smi").capacity -@pytest.mark.xfail( - strict=True, - reason=( - "config.py FIELD_BOUNDS['opt_steps'] declares a floor of ('ge', 1), " - "but Auto3D.utils.validation.check_valid_configuration (validation.py:621) " - "and check_input (validation.py:354) each hand-write a floor of 10. " - "A value FIELD_BOUNDS calls valid is rejected by the other validator. " - "This test does not decide which number is right -- only that the two " - "must agree on ONE minimum, whichever the eventual fix picks." - ), -) -def test_opt_steps_minimum_agrees_between_config_and_validation(): - """FIELD_BOUNDS['opt_steps'] and utils.validation's opt_steps floor must - be the SAME minimum -- see the xfail reason for the defect this pins. +def test_opt_steps_has_exactly_one_declared_minimum(): + """``opt_steps`` had TWO minimums: ``FIELD_BOUNDS`` declared ``("ge", 1)`` + while ``utils/validation.py`` hand-wrote ``< 10`` twice (in ``check_input`` + and again in ``check_valid_configuration``). So ``opt_steps=5`` was accepted + by ``Auto3DOptions``/``CLIConfig``, printed a banner, and only then failed + at run start. + + 10 won, not 1 -- see the comment on ``FIELD_BOUNDS["opt_steps"]`` for why + (the optimizer's own 10-step cadences, and its explicit ``n >= 10`` guard). + This test asserts the *consolidation*, in both directions: the floor is + declared once, in ``FIELD_BOUNDS``, and no validator downstream carries a + second copy of it. """ - from Auto3D.config import FIELD_BOUNDS, check_field_bounds - from Auto3D.utils.validation import check_valid_configuration + from Auto3D.config import FIELD_BOUNDS, Auto3DOptions, check_field_bounds + from Auto3D.exceptions import ConfigurationError + from Auto3D.utils import validation as validation_mod kind, bound_min = FIELD_BOUNDS["opt_steps"] - assert kind == "ge" + assert (kind, bound_min) == ("ge", 10) - # config.py's own gate must accept its own declared floor (premise, not - # the point of this test). + # The declared floor is accepted; one below it is refused -- and refused at + # construction, on the object every entry point builds, not later. check_field_bounds({"opt_steps": bound_min}) # must not raise - - # utils.validation must not disagree with config.py's declared floor: - # a value config.py calls valid must not be flagged as an error there. - errors = check_valid_configuration( - path=None, k=1, opt_steps=int(bound_min), use_gpu=False, - ) - assert not any("opt_steps" in e for e in errors), ( - f"config.py says opt_steps={bound_min} is valid (FIELD_BOUNDS), but " - f"utils.validation.check_valid_configuration disagrees: {errors}" + Auto3DOptions(path="x.smi", k=1, opt_steps=int(bound_min)) + with pytest.raises(ConfigurationError, match="opt_steps"): + Auto3DOptions(path="x.smi", k=1, opt_steps=int(bound_min) - 1) + + # And utils/validation.py no longer restates the number anywhere. A source + # check rather than a behavioral one: a second copy that happens to agree + # today is exactly how the two drifted apart in the first place, and + # behavior cannot distinguish "one bound" from "two bounds that match". + # + # AST, not a substring scan over lines: the prose explaining why the checks + # were removed necessarily mentions `opt_steps` and `< 10`, so a text match + # flags the comment that documents the fix. A `Compare` node is the thing + # actually forbidden here. + import ast + + tree = ast.parse(Path(validation_mod.__file__).read_text()) + offenders = [ + ast.unparse(node) + for node in ast.walk(tree) + if isinstance(node, ast.Compare) + and "opt_steps" in ast.unparse(node) + and any( + isinstance(c, ast.Constant) and isinstance(c.value, (int, float)) + for c in node.comparators + ) + ] + assert not offenders, ( + f"utils/validation.py hand-writes an opt_steps bound again: {offenders}. " + "The bound belongs in Auto3D.config.FIELD_BOUNDS only." ) diff --git a/tests/test_config_parity.py b/tests/test_config_parity.py index d4762fa3..bb1e15dd 100644 --- a/tests/test_config_parity.py +++ b/tests/test_config_parity.py @@ -165,7 +165,7 @@ def forward_batched(self, coords, numbers, charges, atom_mask=None): monkeypatch.setattr(spe_mod, "EnForce_ANI", FakeEnForce) - def fake_pad(mols, model_name, device, coord_pad, species_pad): + def fake_pad(mols, adapter, device): n = len(mols) coords = torch.zeros(n, 1, 3) numbers = torch.zeros(n, 1, dtype=torch.long) @@ -468,7 +468,7 @@ class _FakeOptimizing: executes exactly as it would after a genuine optimization -- without loading an NNP.""" - def __init__(self, in_f, out_f, name, device, config, *a, **k): + def __init__(self, in_f, out_f, *a, **k): self.in_f = in_f self.out_f = out_f diff --git a/tests/test_custom_nnp_contract.py b/tests/test_custom_nnp_contract.py index e6c722a7..6da0ff2a 100644 --- a/tests/test_custom_nnp_contract.py +++ b/tests/test_custom_nnp_contract.py @@ -470,26 +470,44 @@ def forward(self, species, coords, charges): CustomModelAdapter(str(tmp_path / "unused.pt"), CPU) -def test_base_adapter_species_pad_default_agrees_with_the_padding_layer(): - """One default, not two. - - BaseModelAdapter used to default species_pad to 0 while - batch_opt.padding.pad_from_mols defaults it to -1, so a subclass that did - not pass the value got a different notion of padding depending on which - layer supplied it -- and 0 collides with ANI2xt's hydrogen index. -1 wins: - it can be neither an atomic number nor a 0-based species index. +def test_the_adapters_pad_is_what_the_padder_writes(): + """One source, not two agreeing sources. + + This used to compare two independent DEFAULTS -- ``BaseModelAdapter``'s + ``species_pad`` against ``pad_from_mols``'s own -- because each layer had its + own, and they disagreed (0 vs -1, where 0 collides with ANI2xt's hydrogen + index). ``pad_from_mols`` now has no pad parameter at all: it reads the value + off the adapter it is padding for. So the assertion becomes the stronger, + structural one -- whatever the adapter says is what lands in the tensor -- and + the old comparison is not merely satisfied but meaningless. """ import inspect + from rdkit import Chem + from rdkit.Chem import AllChem + from Auto3D.batch_opt.padding import pad_from_mols from Auto3D.models.adapter import BaseModelAdapter + from tests.helpers_adapter import FakeAdapter + + assert "species_pad" not in inspect.signature(pad_from_mols).parameters + assert "coord_pad" not in inspect.signature(pad_from_mols).parameters - adapter_default = inspect.signature( + # -1 remains the safe default for a third-party subclass: it can be neither + # a real atomic number nor a 0-based species index. + assert inspect.signature( BaseModelAdapter.__init__ - ).parameters["species_pad"].default - padding_default = inspect.signature(pad_from_mols).parameters["species_pad"].default + ).parameters["species_pad"].default == -1 - assert adapter_default == padding_default == -1 + mols = [] + for smiles in ("C", "O"): # 5 atoms and 3, so the batch is padded + mol = Chem.AddHs(Chem.MolFromSmiles(smiles)) + assert AllChem.EmbedMolecule(mol, randomSeed=42) == 0 + mols.append(mol) + adapter = FakeAdapter(coord_pad=7.25, species_pad=-99) + coords, species, _charges, _mask = pad_from_mols(mols, adapter, CPU) + assert species[1, 3:].tolist() == [-99, -99] + assert torch.all(coords[1, 3:] == 7.25) def test_validate_custom_nnp_is_callable_directly(): @@ -581,13 +599,11 @@ def _padded_batch(model): mol = Chem.AddHs(Chem.MolFromSmiles(smiles)) assert AllChem.EmbedMolecule(mol, randomSeed=42) == 0 mols.append(mol) - coords, species, charges, _atom_mask = pad_from_mols( - mols, - "AIMNET", - CPU, - coord_pad=model.coord_pad, - species_pad=model.species_pad, - ) + # The example model IS the adapter here: it declares its own + # coord_pad/species_pad and consumes raw atomic numbers, so the identity + # `to_species` is what a custom NNP must get. + model.to_species = list + coords, species, charges, _atom_mask = pad_from_mols(mols, model, CPU) return mols, coords, species, charges @pytest.mark.parametrize("module_name", EXAMPLE_MODULES) @@ -630,3 +646,116 @@ def __call__(self, inputs, forces=False): f"{module_name}.userNNP2 dropped the R-group (Z=0) atom, so it " "scored a different molecule than the one submitted" ) + + +# --- the contract is derived from the Protocol, not retyped next to it ------ + +def test_required_attributes_tracks_the_protocol(): + """What the validator demands must be DERIVED from ``CustomNNP`` itself. + + ``REQUIRED_ATTRIBUTES`` used to be a hand-written tuple sitting a few lines + above the Protocol that declares the same two members, with nothing linking + them. This test enumerates the Protocol *here* rather than hardcoding names, + so adding a data member to ``CustomNNP`` without the validator learning + about it in the same edit goes red. + """ + from Auto3D.models.contract import REQUIRED_ATTRIBUTES, CustomNNP + + declared = tuple(CustomNNP.__annotations__) + assert REQUIRED_ATTRIBUTES == declared, ( + f"validator demands {REQUIRED_ATTRIBUTES} but CustomNNP declares " + f"{declared}; the two must not be maintained separately" + ) + + class NoPadsAtAll: + def forward(self, species, coords, charges): + return coords + + with pytest.raises(ModelLoadError) as excinfo: + validate_custom_nnp(NoPadsAtAll(), "") + message = str(excinfo.value) + for name in declared: + assert name in message, ( + f"{name} is declared on CustomNNP but the rejection message does " + f"not name it: {message}" + ) + + +def test_customnnp_data_members_are_exactly_the_two_padding_values(): + """A deliberate change-detector, not a restatement of the line above. + + ``validate_custom_nnp`` skips the ``forward`` signature check entirely for a + TorchScript ``RecursiveScriptModule`` (its forward is a pybind11 builtin + with no Python signature), so for every archive in the wild + ``REQUIRED_ATTRIBUTES`` is the ONLY gate. Now that the tuple is derived from + ``CustomNNP.__annotations__``, adding an annotated field to the Protocol -- + even "just for documentation" -- immediately rejects every existing archive + that does not carry it. That is a breaking change and must be released as + one; this test is what makes it impossible to do by accident. + """ + from Auto3D.models.contract import CustomNNP + + assert tuple(CustomNNP.__annotations__) == ("coord_pad", "species_pad") + + +def test_customnnp_is_not_runtime_checkable(): + """``isinstance(x, CustomNNP)`` must raise, not answer. + + ``@runtime_checkable`` tests attribute *presence* only. Every + ``torch.nn.Module`` has a ``forward`` attribute (torch installs + ``Module.forward = _forward_unimplemented``), so the single most common real + failure -- a saved module that never defined its own ``forward`` -- would + pass an ``isinstance`` check while raising ``NotImplementedError`` deep in + the optimization loop. A boolean also cannot carry the diagnosis + ``validate_custom_nnp`` produces. So the honest answer to "can I check this + at runtime?" is a ``TypeError`` pointing at the validator. + """ + from Auto3D.models.contract import CustomNNP + + with pytest.raises(TypeError): + isinstance(object(), CustomNNP) # noqa: B015 - the call IS the assertion + + +def test_keyword_only_forward_message_is_not_its_own_demand(): + """The rejection must not render the signature it is asking for. + + The message used to comma-join parameter *names*, dropping ``*``, ``/`` and + defaults, so a keyword-only ``forward(self, *, species, coords, charges)`` + was rejected with "has forward(species, coords, charges) ... Expected + forward(self, species, coords, charges)" -- text that shows the author + nothing wrong. Rendering the signature the way the interpreter does keeps + the marker that actually explains the refusal. + """ + from Auto3D.models.contract import EXPECTED_SIGNATURE + + class KeywordOnly: + coord_pad = 0.0 + species_pad = -1 + + def forward(self, *, species, coords, charges): + return coords + + with pytest.raises(ModelLoadError) as excinfo: + validate_custom_nnp(KeywordOnly(), "") + message = str(excinfo.value) + observed = message.split("Expected")[0] + assert "*" in observed, ( + "the keyword-only marker is what explains the rejection, but the " + f"message renders no '*': {message}" + ) + assert EXPECTED_SIGNATURE in message + + +def test_positional_only_marker_survives_the_transposed_message(): + """The order-check message renders ``/`` too, for the same reason.""" + + class Transposed: + coord_pad = 0.0 + species_pad = -1 + + def forward(self, coords, species, /, charges): + return coords + + with pytest.raises(ModelLoadError) as excinfo: + validate_custom_nnp(Transposed(), "") + assert "/" in str(excinfo.value).split("but Auto3D calls")[0] diff --git a/tests/test_durability.py b/tests/test_durability.py index eb376ea2..07db026e 100644 --- a/tests/test_durability.py +++ b/tests/test_durability.py @@ -207,7 +207,7 @@ class FakeOptimizing: unrecoverable before C14 was fixed. """ - def __init__(self, path, outpath, model_name, device, opt_config): + def __init__(self, path, outpath, *, adapter, device, config): self._path = path self._outpath = outpath @@ -404,7 +404,7 @@ def forward_batched(self, coords, numbers, charges, atom_mask=None): monkeypatch.setattr(spe_mod, "EnForce_ANI", FakeEnForce) - def fake_pad(mols, model_name, device, coord_pad, species_pad): + def fake_pad(mols, adapter, device): n = len(mols) coords = torch.zeros(n, 1, 3) numbers = torch.zeros(n, 1, dtype=torch.long) @@ -453,7 +453,7 @@ def test_opt_geometry_rejects_output_equal_to_input(self, job_dir, monkeypatch): original = sdf.read_bytes() class FakeOptimizing: - def __init__(self, path, outpath, model_name, device, opt_config): + def __init__(self, path, outpath, *, adapter, device, config): self._path = path self._outpath = outpath @@ -865,7 +865,7 @@ def forward_batched(self, coords, numbers, charges, atom_mask=None): n = coords.shape[0] return torch.ones(n, dtype=torch.float64), torch.zeros_like(coords) - def fake_pad(mols, model_name, device, coord_pad, species_pad): + def fake_pad(mols, adapter, device): n = len(mols) return ( torch.zeros(n, 1, 3), diff --git a/tests/test_e_tot_units.py b/tests/test_e_tot_units.py index 62971f4f..e3b591ba 100644 --- a/tests/test_e_tot_units.py +++ b/tests/test_e_tot_units.py @@ -13,15 +13,17 @@ where the truth is 1.000, and the ranker's own eV->Hartree conversion runs on an already-Hartree number, dividing by 27.211 twice. -No neural network potential is loaded: ``ensemble_opt`` and ``create_model`` -are stubbed at the model boundary, exactly as ``tests/test_batchopt.py`` does, -so the real padder, the real ``optimizing.run`` writer, the real -``_annotate_and_rewrite`` and the real ranker all execute. +No neural network potential is loaded: ``ensemble_opt`` is stubbed and the +adapter is a conforming double, so the real padder, the real ``optimizing.run`` +writer, the real ``_annotate_and_rewrite`` and the real ranker all execute. + +``optimizing`` no longer builds its own adapter (audit M41), so the direct tests +inject one and the ``opt_geometry`` tests stub ``create_model`` where +``opt_geometry`` itself reads it -- ``Auto3D.ASE.geometry`` -- rather than at a +seam inside ``batch_opt``. """ from __future__ import annotations -from types import SimpleNamespace - import pytest import torch from rdkit import Chem @@ -51,8 +53,16 @@ def _write_input(path, names) -> None: def _stub_model_boundary(monkeypatch, energies_ev): - """Replace the NNP with a table of energies; everything else stays real.""" + """Replace the NNP with a table of energies; everything else stays real. + + Returns the conforming adapter double, for callers that construct + ``optimizing`` directly. ``Auto3D.ASE.geometry.create_model`` is stubbed to + hand back the same object, because that is where ``opt_geometry`` now builds + the adapter it injects. + """ + import Auto3D.ASE.geometry as geo import Auto3D.batch_opt.batchopt as bo + from tests.helpers_adapter import FakeAdapter def fake_ensemble_opt(net, coord, numbers, charges, param, device, atom_mask=None, progress_cb=None): @@ -64,11 +74,10 @@ def fake_ensemble_opt(net, coord, numbers, charges, param, device, converged_mask=[True] * n, oscillating_count=[0] * n, ) + adapter = FakeAdapter() monkeypatch.setattr(bo, "ensemble_opt", fake_ensemble_opt) - monkeypatch.setattr( - bo, "create_model", - lambda *a, **k: SimpleNamespace(coord_pad=0.0, species_pad=-1), - ) + monkeypatch.setattr(geo, "create_model", lambda *a, **k: adapter) + return adapter class TestOptimizerWritesHartree: @@ -77,14 +86,15 @@ class TestOptimizerWritesHartree: def test_e_tot_is_the_model_energy_in_hartree(self, tmp_path, monkeypatch): import Auto3D.batch_opt.batchopt as bo - _stub_model_boundary(monkeypatch, ENERGIES_EV) + adapter = _stub_model_boundary(monkeypatch, ENERGIES_EV) inp = tmp_path / "in.sdf" out = tmp_path / "out.sdf" _write_input(inp, ["spec_0_0", "spec_0_1", "spec_0_2"]) bo.optimizing( - str(inp), str(out), "AIMNET", torch.device("cpu"), - {"opt_steps": 1, "opttol": 0.01, "patience": 1, "batchsize_atoms": 1024}, + str(inp), str(out), adapter=adapter, device=torch.device("cpu"), + config={"opt_steps": 1, "opttol": 0.01, "patience": 1, + "batchsize_atoms": 1024}, ).run() mols = [m for m in Chem.SDMolSupplier(str(out), removeHs=False) if m] @@ -178,13 +188,14 @@ def test_relative_tautomer_energy_is_kcal_per_mol(self, tmp_path, monkeypatch): import Auto3D.batch_opt.batchopt as bo from Auto3D.tautomer import select_tautomers - _stub_model_boundary(monkeypatch, ENERGIES_EV[:2]) + adapter = _stub_model_boundary(monkeypatch, ENERGIES_EV[:2]) inp = tmp_path / "in.sdf" out = tmp_path / "opt.sdf" _write_input(inp, ["id1@taut0_0_0", "id1@taut1_0_0"]) bo.optimizing( - str(inp), str(out), "AIMNET", torch.device("cpu"), - {"opt_steps": 1, "opttol": 0.01, "patience": 1, "batchsize_atoms": 1024}, + str(inp), str(out), adapter=adapter, device=torch.device("cpu"), + config={"opt_steps": 1, "opttol": 0.01, "patience": 1, + "batchsize_atoms": 1024}, ).run() selected = select_tautomers(str(out), k=2) diff --git a/tests/test_isomer_engine_hardening.py b/tests/test_isomer_engine_hardening.py index f96bd621..ab900535 100644 --- a/tests/test_isomer_engine_hardening.py +++ b/tests/test_isomer_engine_hardening.py @@ -289,7 +289,7 @@ def forward_batched(self, coords, numbers, charges, atom_mask=None): monkeypatch.setattr(spe_mod, "EnForce_ANI", FakeEnForce) - def fake_pad(mols, model_name, device, coord_pad, species_pad): + def fake_pad(mols, adapter, device): assert all(m is not None for m in mols), "None leaked into pad_from_mols" n = len(mols) coords = torch.zeros(n, 1, 3, requires_grad=True) diff --git a/tests/test_legacy_yaml_parity.py b/tests/test_legacy_yaml_parity.py new file mode 100644 index 00000000..235caf8d --- /dev/null +++ b/tests/test_legacy_yaml_parity.py @@ -0,0 +1,144 @@ +"""One YAML ingestion path: the two entry points must agree on malformed files. + +Auto3D has two ways to hand it a YAML configuration: + +* the deprecated ``auto3d `` form (``auto3Dcli._run_legacy_yaml``); +* the modern ``auto3d run INPUT -c `` form + (``cli.commands.run.execute_run`` -> ``cli.config_schema.load_yaml_config``). + +Both already shared every *value* validator -- ``build_cli_config``, and through +it ``FIELD_BOUNDS``, ``extra="forbid"``, ``parse_gpu_idx`` and the engine +registry lookup. What they did **not** share was the *ingestion* layer: the +three shape guards in ``load_yaml_config`` (empty file, non-mapping top level, +unparseable YAML). ``_run_legacy_yaml`` carried its own ``yaml.safe_load`` and +its own ``"None"``-string loop, so an empty or list-topped file reached +``parameters.items()`` and surfaced as ``AttributeError``/``TypeError`` under +the generic "Unexpected Error" panel at **exit 1**, while the identical file +through ``-c`` gave a ``ConfigurationError`` at **exit 2** with a hint. + +Two exit codes for one file is exactly what ``build_cli_config``'s docstring +says it exists to prevent, so this module asserts the property directly rather +than asserting either path's behavior in isolation: for each malformed shape, +both entry points must report the **same exception class** and the **same exit +code**. +""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +# Each shape is a file the user could plausibly write by mistake. The values are +# the literal file contents; see test_config_parity.py for the well-formed +# counterpart of this comparison (same configuration, every entry point). +MALFORMED_SHAPES: dict[str, str] = { + # yaml.safe_load returns None for an empty document. + "empty_file": "", + # A top-level sequence: valid YAML, but not a mapping of option names. + "top_level_list": "- k: 1\n- path: mols.smi\n", + # A top-level scalar -- the shape you get from a file holding only a comment + # line's worth of text, and the other half of the "not a mapping" case. + "top_level_scalar": "k=1\n", + # Genuinely unparseable: yaml.safe_load raises yaml.YAMLError. + "yaml_syntax_error": "k: 1\n window: 5.0\n", +} + + +def _write(tmp_path: Path, name: str, text: str) -> Path: + path = tmp_path / name + path.write_text(text) + return path + + +def _spy_on(monkeypatch, module, sink: list) -> None: + """Record every exception ``module.handle_error`` is handed, then let the + real handler run so the exit code it chooses is the one under test. + + Patched per-module rather than on ``Auto3D.cli.errors`` alone because the + two entry points bind the name differently: ``cli/commands/run.py`` imports + ``handle_error`` at module level (so the binding must be replaced on *that* + module), while ``_run_legacy_yaml`` imports it inside the function body on + every call. + """ + real = module.handle_error + + def spy(error, *args, **kwargs): + sink.append(error) + return real(error, *args, **kwargs) + + monkeypatch.setattr(module, "handle_error", spy) + + +def _verdict_legacy(yaml_path: Path, monkeypatch) -> tuple[str, int]: + """(exception class name, exit code) from ``auto3d ``.""" + from Auto3D.auto3Dcli import _run_legacy_yaml + from Auto3D.cli import errors as errors_mod + + seen: list[BaseException] = [] + _spy_on(monkeypatch, errors_mod, seen) + + with pytest.raises(SystemExit) as exc_info: + _run_legacy_yaml(str(yaml_path)) + + assert seen, "the legacy path exited without routing through handle_error" + return type(seen[-1]).__name__, exc_info.value.code + + +def _verdict_modern(input_file: Path, yaml_path: Path, monkeypatch) -> tuple[str, int]: + """(exception class name, exit code) from ``auto3d run INPUT -c cfg``.""" + from Auto3D.cli.commands import run as run_mod + + seen: list[BaseException] = [] + _spy_on(monkeypatch, run_mod, seen) + + with pytest.raises(SystemExit) as exc_info: + run_mod.execute_run(input_file=input_file, config_file=yaml_path, quiet=True) + + assert seen, "the modern path exited without routing through handle_error" + return type(seen[-1]).__name__, exc_info.value.code + + +@pytest.mark.parametrize("shape", sorted(MALFORMED_SHAPES)) +def test_malformed_yaml_is_judged_identically_by_both_entry_points( + shape, tmp_path, monkeypatch +): + """A malformed config file must produce the same error class and the same + exit code whichever entry point reads it. + + Before the fix all four shapes gave ``ConfigurationError``/exit 2 through + ``-c`` and an internal-looking ``AttributeError``/``TypeError``/ + ``yaml.YAMLError`` at exit 1 through ``auto3d ``. + """ + cfg = _write(tmp_path, "cfg.yaml", MALFORMED_SHAPES[shape]) + smi = _write(tmp_path, "mols.smi", "CCO m1\n") + + legacy = _verdict_legacy(cfg, monkeypatch) + modern = _verdict_modern(smi, cfg, monkeypatch) + + assert legacy == modern, ( + f"{shape}: legacy 'auto3d cfg.yaml' reported {legacy} but modern " + f"'auto3d run in.smi -c cfg.yaml' reported {modern}" + ) + # Pin the shared verdict too, not just the agreement: a future change that + # made *both* paths crash with "Unexpected Error" at exit 1 would satisfy + # the equality above while destroying the property it exists to protect. + assert modern == ("ConfigurationError", 2), modern + + +def test_settings_only_config_still_refused_by_the_legacy_form(tmp_path, monkeypatch): + """Behavior that must be *preserved*, not changed, by the ingestion merge. + + ``CLIConfig.path`` is optional so a settings-only config file is reusable + across runs (``auto3d run INPUT -c cfg.yaml`` supplies the input). The + deprecated form has no other source of an input path, so + ``to_auto3d_options()`` must still refuse it -- as a ``ConfigurationError`` + at exit 2, naming the missing key. + """ + # ANI2xt, not the default AIMNET: resolving the engine name for a *valid* + # config actually happens here, and AIMNET would import the optional + # `aimnet` package (see tests/test_config_parity.py's matching note). + cfg = _write( + tmp_path, "settings_only.yaml", "k: 1\noptimizing_engine: 'ANI2xt'\n" + ) + + assert _verdict_legacy(cfg, monkeypatch) == ("ConfigurationError", 2) diff --git a/tests/test_model_adapter.py b/tests/test_model_adapter.py index 0777b850..0ea43bb9 100644 --- a/tests/test_model_adapter.py +++ b/tests/test_model_adapter.py @@ -45,14 +45,57 @@ def test_model_adapter_interface(aimnet_model): class TestModelAdapterProtocol: - """Tests for the ModelAdapter protocol.""" + """Tests for the ModelAdapter protocol. - def test_protocol_is_defined(self): - """Protocol should be properly defined.""" - from Auto3D.models.adapter import ModelAdapter + It lives in ``Auto3D.models.contract`` alongside ``CustomNNP`` -- the + interface it is constantly confused with -- so the two declarations cannot be + read separately. ``Auto3D.models.adapter`` holds implementations only. + """ + + def test_protocol_declares_the_whole_contract(self): + from Auto3D.models.contract import ModelAdapter + + for member in ("forward", "energy", "to_species"): + assert hasattr(ModelAdapter, member) + assert tuple(ModelAdapter.__annotations__) == ("coord_pad", "species_pad") + + def test_protocol_does_not_live_in_the_implementation_module(self): + """A clean sweep: the old import path is gone, not aliased.""" + import Auto3D.models.adapter as adapter_mod + + assert not hasattr(adapter_mod, "ModelAdapter") + + def test_the_package_still_re_exports_it(self): + from Auto3D.models import ModelAdapter as reexported + from Auto3D.models.contract import ModelAdapter as canonical + + assert reexported is canonical + assert "ModelAdapter" in __import__( + "Auto3D.models", fromlist=["__all__"] + ).__all__ + + def test_device_is_not_part_of_the_contract(self): + """Dropped deliberately. + + Nothing outside an adapter reads ``adapter.device`` + (``BaseModelAdapter`` keeps ``self.device`` as an implementation detail), + and requiring it would make every legitimate structural implementation -- + including test doubles that never touch a device -- non-conforming for no + benefit. + """ + from Auto3D.models.contract import ModelAdapter + + assert "device" not in ModelAdapter.__annotations__ + + def test_issubclass_is_never_a_valid_question(self): + """Pinned so nobody "improves" the EnForce_ANI gate into an issubclass: + a Protocol with data members raises for it.""" + import pytest - # Check that ModelAdapter has forward method - assert hasattr(ModelAdapter, 'forward') + from Auto3D.models.contract import ModelAdapter + + with pytest.raises(TypeError): + issubclass(dict, ModelAdapter) class TestBaseModelAdapter: @@ -431,3 +474,241 @@ def test_inf_forces_raises(self): forces[0, 2, 1] = float("inf") with pytest.raises(NumericalError, match="Inf.*force"): _validate_outputs(energy, forces) + + +class TestAni2xtNetworksAreTableDriven: + """The seven per-element MLPs are one factory + a width table (audit M61). + + They used to be seven copy-pasted ``nn.Sequential`` blocks, 69 lines + differing only in three integers, which is how ``F_network`` and + ``S_network`` ended up declared in the opposite order from the + ``ModuleList`` they are placed into -- readable only by cross-checking two + distant lines. + + The refactor is safe because ``nn.ModuleList.load_state_dict`` keys off + POSITION (``"0.0.weight"``, ``"1.0.weight"``, ...), never off the Python + variable a submodule was assigned to. That claim is not taken on trust: the + test below loads the real shipped checkpoint into both the old hand-written + structure and the new generated one and compares every tensor. + """ + + CHECKPOINT = "src/Auto3D/models/ani2xt_no_repulsion.pt" + + @staticmethod + def _hand_written(aev_dim): + """Verbatim transcription of the seven blocks as they were before M61, + in their original ``ModuleList`` order (note F before S).""" + from torch import nn + + def mlp(a, b, c, d): + return nn.Sequential( + nn.Linear(a, b), nn.CELU(0.1), + nn.Linear(b, c), nn.CELU(0.1), + nn.Linear(c, d), nn.CELU(0.1), + nn.Linear(d, 1), + ) + + H_network = mlp(aev_dim, 256, 192, 160) + C_network = mlp(aev_dim, 224, 192, 160) + N_network = mlp(aev_dim, 192, 160, 128) + O_network = mlp(aev_dim, 192, 160, 128) + S_network = mlp(aev_dim, 160, 128, 96) + F_network = mlp(aev_dim, 160, 128, 96) + Cl_network = mlp(aev_dim, 160, 128, 96) + return nn.ModuleList([ + H_network, C_network, N_network, O_network, + F_network, S_network, Cl_network, + ]) + + def test_the_shipped_checkpoint_loads_identically_both_ways(self): + """The tripwire for the whole refactor: same weights, tensor for tensor. + + Loading a ``state_dict`` is not running the model -- no inference, no + torchani, no download; the checkpoint is bundled in + ``src/Auto3D/models/``. + """ + import torch + from torch import nn + + from Auto3D.batch_opt.ANI2xt_no_rep import WIDTHS, _atomic_mlp + + checkpoint = torch.load(self.CHECKPOINT, map_location="cpu", weights_only=True) + # Read the AEV width off the checkpoint rather than hardcoding it, so a + # retrained model with a different AEV does not silently pass. + aev_dim = checkpoint["0.0.weight"].shape[1] + + old = self._hand_written(aev_dim) + new = nn.ModuleList([_atomic_mlp(aev_dim, widths) for widths in WIDTHS]) + + old.load_state_dict(checkpoint) + new.load_state_dict(checkpoint) + + old_state, new_state = old.state_dict(), new.state_dict() + assert set(old_state) == set(new_state) + for key in old_state: + assert torch.equal(old_state[key], new_state[key]), ( + f"{key} differs: the ModuleList order changed, so this " + f"checkpoint now routes an element to the wrong network" + ) + + def test_the_width_table_is_in_moduleList_order_including_f_before_s(self): + """Fluorine at index 4 and sulfur at index 5, matching ANI2XT_INDEX. + + The old code declared ``S_network`` before ``F_network`` in the source but + placed F before S in the ``ModuleList``. Since F, S and Cl happen to share + the same widths the mix-up was harmless -- and undetectable. The table + makes the order the only order there is. + """ + from Auto3D.batch_opt.ANI2xt_no_rep import WIDTHS + from Auto3D.models.species import ANI2XT_INDEX + + assert len(WIDTHS) == len(ANI2XT_INDEX) == 7 + # H, C, N, O, F, S, Cl + assert WIDTHS == ( + (256, 192, 160), + (224, 192, 160), + (192, 160, 128), + (192, 160, 128), + (160, 128, 96), + (160, 128, 96), + (160, 128, 96), + ) + + def test_the_factory_builds_the_documented_shape(self): + from torch import nn + + from Auto3D.batch_opt.ANI2xt_no_rep import _atomic_mlp + + net = _atomic_mlp(11, (5, 4, 3)) + kinds = [type(layer) for layer in net] + assert kinds == [ + nn.Linear, nn.CELU, nn.Linear, nn.CELU, nn.Linear, nn.CELU, nn.Linear + ] + assert [ (l.in_features, l.out_features) + for l in net if isinstance(l, nn.Linear) ] == [ + (11, 5), (5, 4), (4, 3), (3, 1) + ] + assert all(l.alpha == 0.1 for l in net if isinstance(l, nn.CELU)) + + +class TestEnergyIsDtypePreserving: + """``energy()`` must answer in the dtype it was asked in. + + This is the single most likely silent numerical regression in the whole + contract change, and it produces no error of any kind. + ``ANI2xAdapter.forward`` and ``CustomModelAdapter.forward`` both call + ``coords.float()`` -- correct for them, because they front float32 weights. + But ``energy()`` exists so a caller can DIFFERENTIATE it (an fp64 Hessian, + which ``ASE/thermo.py`` builds by promoting the wrapped module with + ``.double()``). If ``energy`` were the inherited ``forward(...)[0]`` for those + two adapters, an fp64 request would come back fp32 with nothing reported: the + Hessian would simply be less accurate than the code around it promises. + + Each adapter is built by calling ``BaseModelAdapter.__init__`` on a bypassed + instance and handing it a toy module that RECORDS the dtype it was fed -- the + same technique the force-sign tests above use. Nothing is loaded, no torchani, + no download. + """ + + @staticmethod + def _bypassed(cls, model): + from Auto3D.models.adapter import BaseModelAdapter + + adapter = cls.__new__(cls) + BaseModelAdapter.__init__( + adapter, model, torch.device("cpu"), coord_pad=0.0, species_pad=-1 + ) + return adapter + + def test_ani2x_energy_keeps_float64(self): + from collections import namedtuple + + from Auto3D.constants import HARTREE_TO_EV + from Auto3D.models.adapter import ANI2xAdapter + + _SpeciesEnergies = namedtuple("SpeciesEnergies", ["species", "energies"]) + seen: list = [] + + class _Recorder(torch.nn.Module): + def forward(self, species_coords): + species, coords = species_coords + seen.append(coords.dtype) + return _SpeciesEnergies( + species, (coords ** 2).sum(dim=(1, 2)) / HARTREE_TO_EV + ) + + adapter = self._bypassed(ANI2xAdapter, _Recorder()) + coords = torch.randn(2, 4, 3, dtype=torch.float64) + species = torch.tensor([[1, 6, 7, 8], [1, 6, 7, -1]]) + charges = torch.zeros(2, dtype=torch.float64) + + energy = adapter.energy(coords, species, charges) + assert seen == [torch.float64], ( + f"energy() handed the model {seen}; an fp64 caller silently got fp32" + ) + assert energy.dtype is torch.float64 + + # And forward() still downcasts, deliberately: that is the optimization + # path, where float32 weights are the point. + seen.clear() + adapter.forward(coords.clone(), species, charges) + assert seen == [torch.float32] + + def test_custom_energy_keeps_float64_for_coords_and_charges(self): + from Auto3D.models.adapter import CustomModelAdapter + + seen: list = [] + + class _Recorder(torch.nn.Module): + def forward(self, species, coords, charges): + seen.append((coords.dtype, charges.dtype)) + return (coords ** 2).sum(dim=(1, 2)) + + adapter = self._bypassed(CustomModelAdapter, _Recorder()) + coords = torch.randn(2, 4, 3, dtype=torch.float64) + species = torch.tensor([[1, 6, 7, 8], [1, 6, 7, -1]]) + charges = torch.zeros(2) + + energy = adapter.energy(coords, species, charges) + assert seen == [(torch.float64, torch.float64)], ( + f"energy() handed the model {seen}; charges must follow coords so a " + "model that concatenates them does not hit a dtype mismatch" + ) + assert energy.dtype is torch.float64 + + seen.clear() + adapter.forward(coords.clone(), species, charges) + assert seen == [(torch.float32, torch.float32)] + + def test_ani2xt_energy_accepts_a_non_leaf_tensor(self): + """``ANI2xtAdapter.forward`` calls ``coords.requires_grad_(True)``, which + raises on the non-leaf tensor an autograd Hessian hands in. Its own + ``energy`` must not touch ``requires_grad`` at all.""" + from Auto3D.models.adapter import ANI2xtAdapter + + class _Toy(torch.nn.Module): + def forward(self, species, coords): + return (coords ** 2).sum(dim=(1, 2)) + + adapter = self._bypassed(ANI2xtAdapter, _Toy()) + leaf = torch.randn(2, 4, 3, dtype=torch.float64, requires_grad=True) + non_leaf = leaf * 2.0 + assert non_leaf.grad_fn is not None, "test premise: coords must be non-leaf" + species = torch.tensor([[0, 1, 2, 3], [0, 1, 2, -1]]) + + energy = adapter.energy(non_leaf, species, torch.zeros(2)) + assert energy.dtype is torch.float64 + # Still graph-connected: no internal no_grad, so a Hessian can be taken. + assert energy.requires_grad + (grad,) = torch.autograd.grad(energy.sum(), leaf) + torch.testing.assert_close(grad, 8.0 * leaf) + + def test_energy_has_no_no_grad_anywhere(self): + """The contract promises a graph-connected result for every adapter.""" + from tests.helpers_adapter import FakeAdapter + + adapter = FakeAdapter() + coords = torch.randn(1, 3, 3, requires_grad=True) + energy = adapter.energy(coords, torch.ones(1, 3, dtype=torch.long), + torch.zeros(1)) + assert energy.requires_grad diff --git a/tests/test_model_factory.py b/tests/test_model_factory.py index c99554a2..7af573b8 100644 --- a/tests/test_model_factory.py +++ b/tests/test_model_factory.py @@ -13,7 +13,39 @@ get_device, is_custom_model, ) -from Auto3D.models.adapter import ModelAdapter +from Auto3D.models.contract import ModelAdapter + + +def test_the_factory_promises_the_contract_not_the_base_class(): + """Every factory signature must be annotated with the Protocol. + + ``ModelAdapter`` was ``@runtime_checkable`` and published, while every + signature that wanted "an adapter" said ``BaseModelAdapter`` (the ABC) + instead -- so the contract the factory actually honors was invisible in its + own types, and production quietly accepted structural implementations the + annotation excluded. Reading ``__annotations__`` (strings, because the module + uses ``from __future__ import annotations``) is the only way to observe it. + """ + import inspect + + from Auto3D import model_factory + + assert ( + inspect.get_annotations(model_factory.create_model)["return"] + == "ModelAdapter" + ) + assert ( + inspect.get_annotations(ModelFactory.create.__func__)["return"] + == "ModelAdapter" + ) + # BaseModelAdapter survives in exactly one type position: a registry of + # Auto3D's OWN adapter classes, which really is the concrete base. + assert ( + inspect.get_annotations(ModelFactory)["_adapters"] + == "dict[str, type[BaseModelAdapter]]" + ) + # ...and it is no longer an incidental runtime re-export of this module. + assert not hasattr(model_factory, "BaseModelAdapter") class TestModelFactory: diff --git a/tests/test_model_preflight.py b/tests/test_model_preflight.py index afdc47ba..09d48ba8 100644 --- a/tests/test_model_preflight.py +++ b/tests/test_model_preflight.py @@ -10,13 +10,14 @@ Verified against production (not assumed from the plan): -- ``check_valid_configuration`` (utils/validation.py:267-361) takes explicit - keyword parameters and returns a ``list[str]`` of error messages -- it never - receives an ``Auto3DOptions`` instance and never raises. The raise happens - one layer up, in ``WorkflowOrchestrator._validate_input`` - (workflow.py:164-179), which is what these tests exercise directly. -- ``check_input`` (utils/validation.py:35) never constructs any model adapter - -- it only checks installed dependencies, opt_steps, and input-file format. +- ``check_valid_configuration`` takes an ``Auto3DOptions`` and returns a + ``list[str]`` of error messages; it does not raise for a bad value. The raise + happens one layer up, in ``WorkflowOrchestrator._validate_input``, which is + what these tests exercise directly. (It used to take ten explicit keyword + parameters with their own defaults -- a third configuration schema -- and + received no options object at all.) +- ``check_input`` never constructs any model adapter -- it only checks + installed dependencies and input-file format. Patching ``AIMNet2Calculator`` and calling ``check_input`` (as an earlier draft of this suite assumed) would never intercept anything; the real construction site is ``optimizing.__init__`` (batch_opt/batchopt.py:175), diff --git a/tests/test_model_wrapper.py b/tests/test_model_wrapper.py index 5eb03f1f..9490086f 100644 --- a/tests/test_model_wrapper.py +++ b/tests/test_model_wrapper.py @@ -210,9 +210,16 @@ def test_forward_batched_retries_on_oom(): from Auto3D.batch_opt.model_wrapper import EnForce_ANI - class _OOMAdapter: - coord_pad = 0.0 - species_pad = -1 + from tests.helpers_adapter import FakeAdapter + + class _OOMAdapter(FakeAdapter): + """Conforms to the contract (inherited), then OOMs on purpose. + + Subclassing the shared double rather than re-declaring an ad-hoc one is + what keeps ``EnForce_ANI``'s contract gate tightenable: a hand-rolled + stub listing only the members this test happens to exercise goes red for + a reason that has nothing to do with OOM retry. + """ def forward(self, coord, numbers, charges, atom_mask=None): if coord.shape[0] > 1: @@ -247,3 +254,68 @@ def test_a_model_name_in_the_batchsize_slot_is_rejected(): with pytest.raises(TypeError, match="batchsize_atoms"): EnForce_ANI(MagicMock(), "AIMNET") + + +class TestEnForceANIRejectsNonAdapters: + """The adapter contract is enforced here, at the one seam that consumes it. + + ``ModelAdapter`` was declared ``@runtime_checkable`` and then never checked + anywhere in ``src/`` or ``tests/``, while every signature that wanted "an + adapter" annotated the ABC instead. This class is what makes the Protocol + load-bearing: a category error (a raw ``nn.Module``, an + ``AIMNet2Calculator``, an engine-name string) is named here instead of + surfacing as an ``AttributeError`` several frames deep inside + ``forward_batched``. + + Note what this does NOT catch: presence is not arity. An object with a + wrong-signature ``forward`` still passes, and a ``MagicMock`` passes + trivially (which the tests above rely on). The gate is for category errors. + """ + + def test_a_raw_nn_module_is_rejected_and_the_gap_is_named(self): + import pytest + + with pytest.raises(TypeError) as excinfo: + EnForce_ANI(torch.nn.Linear(1, 1)) + message = str(excinfo.value) + assert "ModelAdapter" in message + # The missing members must be enumerated, not merely alluded to. + for name in ("to_species", "coord_pad", "species_pad", "energy"): + assert name in message, f"{name} is missing but unnamed: {message}" + + def test_an_engine_name_string_is_rejected(self): + """The pre-adapter API took a model name here; a stale caller must not + get an object whose ``forward`` fails much later.""" + import pytest + + with pytest.raises(TypeError, match="ModelAdapter"): + EnForce_ANI("AIMNET") + + def test_a_conforming_double_is_accepted(self): + """The gate must not reject a structural (non-subclass) adapter -- + production has always accepted those, which is why annotating the ABC + instead of the Protocol made the Protocol decorative.""" + from tests.helpers_adapter import FakeAdapter, padded_batch + + adapter = FakeAdapter() + wrapper = EnForce_ANI(adapter) + coords, species, charges, atom_mask = padded_batch() + e, f = wrapper.forward(coords, species, charges, atom_mask=atom_mask) + assert e.shape == (2,) + assert f.shape == (2, 3, 3) + + def test_the_missing_member_list_comes_from_the_protocol(self): + """Derived, not hand-listed: widening ``ModelAdapter`` must widen this + message in the same edit.""" + import pytest + + from Auto3D.models.contract import ModelAdapter + + class NothingAtAll: + pass + + with pytest.raises(TypeError) as excinfo: + EnForce_ANI(NothingAtAll()) + message = str(excinfo.value) + for name in ModelAdapter.__annotations__: + assert name in message diff --git a/tests/test_padding.py b/tests/test_padding.py index b746ac5b..fb859fb6 100644 --- a/tests/test_padding.py +++ b/tests/test_padding.py @@ -6,6 +6,28 @@ from rdkit.Chem import AllChem from Auto3D.batch_opt.padding import pad_from_mols +from tests.helpers_adapter import FakeAdapter + +# The padder now takes the ADAPTER, which supplies the species convention and +# both fill values. These two stand in for the real engines' conventions without +# loading anything: AIMNet2 (raw atomic numbers, species_pad=0) and ANI2xt +# (0-based network indices, species_pad=-1). Constructing the real ANI2xtAdapter +# would load ~7 MB of weights and require torchani for the AEV computer, neither +# of which belongs in the fast tier. +ANI2XT_MAP = {1: 0, 6: 1, 7: 2, 8: 3, 9: 4, 16: 5, 17: 6} + + +def _aimnet_like(species_pad: int = 0) -> FakeAdapter: + return FakeAdapter(coord_pad=0.0, species_pad=species_pad) + + +def _ani2xt_like(species_pad: int = -1) -> FakeAdapter: + """Real ANI2xt remap, including its named ValueError for an out-of-set Z.""" + from Auto3D.models.species import to_ani2xt_species + + adapter = FakeAdapter(coord_pad=0.0, species_pad=species_pad) + adapter.to_species = to_ani2xt_species + return adapter class TestPadFromMols: @@ -23,7 +45,7 @@ def test_basic_rdkit_molecules(self): mols = [mol1, mol2] device = torch.device("cpu") - c, s, q, mask = pad_from_mols(mols, "AIMNET", device, coord_pad=0.0, species_pad=0) + c, s, q, mask = pad_from_mols(mols, _aimnet_like(), device) # Methane has 5 atoms (1C + 4H), water has 3 atoms (1O + 2H) assert c.shape == (2, 5, 3) # max_atoms = 5 @@ -40,7 +62,7 @@ def test_species_values_aimnet(self): mols = [mol] device = torch.device("cpu") - c, s, q, mask = pad_from_mols(mols, "AIMNET", device, coord_pad=0.0, species_pad=0) + c, s, q, mask = pad_from_mols(mols, _aimnet_like(), device) # Carbon is atomic number 6, Hydrogen is 1 species_list = s[0].tolist() @@ -55,7 +77,7 @@ def test_species_values_ani2xt(self): mols = [mol] device = torch.device("cpu") - c, s, q, mask = pad_from_mols(mols, "ANI2xt", device, coord_pad=0.0, species_pad=-1) + c, s, q, mask = pad_from_mols(mols, _ani2xt_like(), device) # ANI2xt mapping: H->0, C->1, N->2, O->3, F->4, S->5, Cl->6 species_list = s[0].tolist() @@ -74,7 +96,7 @@ def test_charges_extracted(self): mols = [mol1, mol2] device = torch.device("cpu") - c, s, q, mask = pad_from_mols(mols, "AIMNET", device, coord_pad=0.0, species_pad=0) + c, s, q, mask = pad_from_mols(mols, _aimnet_like(), device) assert q[0].item() == 0 # Methane is neutral assert q[1].item() == -1 # Hydroxide has -1 charge @@ -85,8 +107,7 @@ def test_charges_are_float(self): AllChem.EmbedMolecule(mol, randomSeed=42) device = torch.device("cpu") - _, _, q_mols, _ = pad_from_mols([mol], "AIMNET", device, - coord_pad=0.0, species_pad=0) + _, _, q_mols, _ = pad_from_mols([mol], _aimnet_like(), device) assert q_mols.dtype == torch.float32 def test_coords_match_conformer(self): @@ -97,7 +118,7 @@ def test_coords_match_conformer(self): mols = [mol] device = torch.device("cpu") - c, s, q, mask = pad_from_mols(mols, "AIMNET", device, coord_pad=0.0, species_pad=0) + c, s, q, mask = pad_from_mols(mols, _aimnet_like(), device) # Get positions from RDKit conf = mol.GetConformer() @@ -117,7 +138,7 @@ def test_requires_grad_enabled(self): mols = [mol] device = torch.device("cpu") - c, s, q, mask = pad_from_mols(mols, "AIMNET", device, coord_pad=0.0, species_pad=0) + c, s, q, mask = pad_from_mols(mols, _aimnet_like(), device) assert c.requires_grad is True @@ -133,7 +154,7 @@ def test_ani2xt_unsupported_element_raises_valueerror(self): device = torch.device("cpu") with pytest.raises(ValueError) as exc: - pad_from_mols(mols, "ANI2xt", device, coord_pad=0.0, species_pad=-1) + pad_from_mols(mols, _ani2xt_like(), device) msg = str(exc.value) assert "ANI2xt" in msg and ("15" in msg or "P" in msg) @@ -161,7 +182,7 @@ def _mol(smiles): return m small, large = _mol("CCO"), _mol("c1ccccc1CCCCO") - _, _, _, atom_mask = pad_from_mols([small, large], "AIMNET", device) + _, _, _, atom_mask = pad_from_mols([small, large], _aimnet_like(), device) assert atom_mask.shape == (2, large.GetNumAtoms()) assert atom_mask[0].sum().item() == small.GetNumAtoms() @@ -185,7 +206,7 @@ def _mol(smiles): # species_pad=0 collides with ANI2xt's hydrogen index. The mask must be # derived from atom counts, so the collision cannot matter. _, species, _, atom_mask = pad_from_mols( - [small, large], "ANI2xt", device, coord_pad=0.0, species_pad=0 + [small, large], _ani2xt_like(species_pad=0), device ) n_small = small.GetNumAtoms() @@ -198,3 +219,52 @@ def _mol(smiles): "sanity check: hydrogens really do sit at species index 0, so a " "value-derived mask would have zeroed them" ) + + +class TestThePadderCannotDisagreeWithTheAdapter: + """One object supplies the remap AND both sentinels, so they cannot conflict. + + ``pad_from_mols`` used to take a model-name *string* plus the adapter's two + pad values as separate arguments (``SPE.py`` and ``batchopt.py`` each passed + all three). The species convention therefore came from one source and the + padding sentinel from another, and nothing structurally prevented them from + contradicting each other -- the shape of audit findings C3/C4. The signature + now takes the adapter, so there is only one source. + """ + + @staticmethod + def _molecules(): + mol1 = Chem.AddHs(Chem.MolFromSmiles("C")) # 5 atoms + AllChem.EmbedMolecule(mol1, randomSeed=42) + mol2 = Chem.AddHs(Chem.MolFromSmiles("O")) # 3 atoms + AllChem.EmbedMolecule(mol2, randomSeed=42) + return [mol1, mol2] + + def test_both_the_remap_and_the_pad_come_from_the_one_object(self): + """A fake declaring a sentinel remap and a distinctive pad must see both + land in the same tensor. Impossible to state before the signature change: + the name decided the remap and the adapter decided the pad.""" + from tests.helpers_adapter import FakeAdapter + + # Deliberately not any real engine's convention: H -> 77, C -> 88, O -> 99. + adapter = FakeAdapter( + coord_pad=-5.5, species_pad=-42, species_map={1: 77, 6: 88, 8: 99} + ) + coords, species, charges, mask = pad_from_mols( + self._molecules(), adapter, torch.device("cpu") + ) + + # Methane: C then 4 H, remapped by the ADAPTER. + assert species[0].tolist() == [88, 77, 77, 77, 77] + # Water: O, H, H remapped; the two padded slots hold the ADAPTER's pad. + assert species[1].tolist() == [99, 77, 77, -42, -42] + assert torch.all(coords[1, 3:] == -5.5) + assert mask[1].tolist() == [True, True, True, False, False] + + def test_the_signature_no_longer_accepts_a_name_or_loose_pads(self): + """A stale caller must fail at the call, not bind a string into the slot + that decides which atoms are padding.""" + import inspect + + parameters = list(inspect.signature(pad_from_mols).parameters) + assert parameters == ["mols", "adapter", "device"], parameters diff --git a/tests/test_padding_invariance.py b/tests/test_padding_invariance.py index c4640788..a96eb50b 100644 --- a/tests/test_padding_invariance.py +++ b/tests/test_padding_invariance.py @@ -50,12 +50,11 @@ def test_energy_unchanged_when_padded(self, engine, atol, device): model = create_model(engine, device) small, large = _mol("CCO"), _mol("c1ccccc1CCCCO") - # Drive the padding convention from the adapter under test, not a - # hardcoded per-engine constant, so the two can never drift apart: - # AIMNet2Adapter uses coord_pad=0.0/species_pad=0 while - # ANI2xtAdapter uses coord_pad=0.0/species_pad=-1 - # (src/Auto3D/models/adapter.py:240, :302). - coord_pad, species_pad = model.coord_pad, model.species_pad + # The padding convention comes from the adapter under test, and it is no + # longer possible for it to come from anywhere else: `pad_from_mols` reads + # the species remap and BOTH fill values off the one object it is handed. + # AIMNet2Adapter uses coord_pad=0.0/species_pad=0 while ANI2xtAdapter uses + # coord_pad=0.0/species_pad=-1. # Alone: no padding at all. The explicit atom_mask is forwarded in # both calls: an adapter that has to strip padding (AIMNet2) takes it @@ -63,15 +62,11 @@ def test_energy_unchanged_when_padded(self, engine, atol, device): # which deletes a real atomic number 0 along with the padding # (audit C13). Without it a padded AIMNET batch reaches the model with # Z=0 ghosts at the origin and returns NaN. - c1, s1, q1, m1 = pad_from_mols( - [small], engine, device, coord_pad=coord_pad, species_pad=species_pad - ) + c1, s1, q1, m1 = pad_from_mols([small], model, device) e_alone = model.forward(c1, s1, q1, atom_mask=m1)[0][0] # Batched with a larger molecule: `small` is now padded to `large`'s size. - c2, s2, q2, m2 = pad_from_mols( - [small, large], engine, device, coord_pad=coord_pad, species_pad=species_pad - ) + c2, s2, q2, m2 = pad_from_mols([small, large], model, device) e_padded = model.forward(c2, s2, q2, atom_mask=m2)[0][0] delta = abs(float(e_alone) - float(e_padded)) diff --git a/tests/test_species_conversion.py b/tests/test_species_conversion.py index 9dcf2911..ffe8d228 100644 --- a/tests/test_species_conversion.py +++ b/tests/test_species_conversion.py @@ -5,8 +5,11 @@ `ASE/thermo.py` and `cli/commands/models.py` previously passed raw atomic numbers instead of converting them, evaluating hydrogen with the carbon network and carbon with the chlorine network (audit C3, C4). Both call sites -now convert through `Auto3D.batch_opt.species.to_model_species` before -calling the model, and the tests below guard against that regressing. +now convert through the MODEL: the batch path asks +`ModelAdapter.to_species` (so the remap cannot disagree with the padding +sentinel that comes from the same object), and the thermo path still goes through +the name-keyed `Auto3D.models.species.to_model_species` residual. The tests below +guard against either regressing. The decisive asymmetry: ANI2x gets periodic_table_index=True at both of its sites (thermo.py:338, models/adapter.py:346), so it was always correct and @@ -39,7 +42,11 @@ def test_pad_from_mols_emits_indices_for_ani2xt(self, device): mol = Chem.AddHs(Chem.MolFromSmiles("C")) AllChem.EmbedMolecule(mol, randomSeed=42) - _, species, _, _ = pad_from_mols([mol], "ANI2xt", device) + from Auto3D.model_factory import create_model + + _, species, _, _ = pad_from_mols( + [mol], create_model("ANI2xt", device), device + ) values = sorted(int(v) for v in species[0]) # ANI2XT_INDEX: H=0, C=1. Methane is one carbon and four hydrogens. @@ -66,7 +73,7 @@ def test_thermo_and_batch_paths_agree_on_methane(self, device): model = create_model("ANI2xt", device) - coords_b, species_b, charges_b, _ = pad_from_mols([mol], "ANI2xt", device) + coords_b, species_b, charges_b, _ = pad_from_mols([mol], model, device) e_batch = float(model.forward(coords_b, species_b, charges_b)[0][0]) thermo_in = mol2aimnet_input(mol, device, model_name="ANI2xt") @@ -117,7 +124,7 @@ def test_health_check_energy_matches_real_methane(self, device): mol = Chem.AddHs(Chem.MolFromSmiles("C")) AllChem.EmbedMolecule(mol, randomSeed=42) model = create_model("ANI2xt", device) - coords, species, charges, _ = pad_from_mols([mol], "ANI2xt", device) + coords, species, charges, _ = pad_from_mols([mol], model, device) reference = float(model.forward(coords, species, charges)[0][0]) # _health_check_energy does not exist -- cli/commands/models.py builds diff --git a/tests/test_species_module.py b/tests/test_species_module.py index 561d78b2..0b932bbd 100644 --- a/tests/test_species_module.py +++ b/tests/test_species_module.py @@ -4,12 +4,19 @@ expects 0-based indices (H=0..Cl=6), not atomic numbers. Before this module existed the conversion was duplicated in three places and omitted in two more (audit C3, C4). + +The table now lives in ``Auto3D.models.species``, not ``Auto3D.batch_opt.species``: +the species convention is a property of the MODEL, and hosting it under +``batch_opt`` made the optimizer package a shared-utility provider for ``ASE/``, +``cli/`` and ``models/``'s own padder. The canonical way to reach it is +``ModelAdapter.to_species`` -- asking the object that also supplies +``species_pad``, so the two cannot disagree. """ from __future__ import annotations import pytest -from Auto3D.batch_opt.species import ANI2XT_INDEX, to_model_species +from Auto3D.models.species import ANI2XT_INDEX, to_model_species class TestAni2xtMapping: @@ -70,3 +77,89 @@ def test_map_contents(self): def test_no_duplicate_indices(self): """Every element gets a distinct network slot.""" assert len(set(ANI2XT_INDEX.values())) == len(ANI2XT_INDEX) + + +class TestOnlyAni2xtRemaps: + """A remap leaking onto any other adapter is a silent, untestable disaster. + + ``CustomModelAdapter`` in particular MUST inherit the identity: a custom NNP + receives atomic numbers and declares its own ``species_pad``, so remapping + for it would feed every third-party model different species than its author + tested against -- on the one path with no in-tree test molecules. This is + audit C3/C4 inverted. + """ + + def test_base_adapter_to_species_is_the_identity(self): + from Auto3D.models.adapter import BaseModelAdapter + + assert BaseModelAdapter.to_species(object(), [1, 6, 17]) == [1, 6, 17] + + def test_ani2xt_is_the_only_adapter_that_overrides_it(self): + from Auto3D.models.adapter import ( + AIMNet2Adapter, + ANI2xAdapter, + ANI2xtAdapter, + BaseModelAdapter, + CustomModelAdapter, + ) + + overriding = { + cls.__name__ + for cls in ( + AIMNet2Adapter, + ANI2xAdapter, + ANI2xtAdapter, + CustomModelAdapter, + ) + if cls.to_species is not BaseModelAdapter.to_species + } + assert overriding == {"ANI2xtAdapter"}, ( + f"{overriding - {'ANI2xtAdapter'}} redefine to_species; every engine " + "but ANI2xt consumes raw atomic numbers" + ) + + def test_ani2xt_adapter_delegates_to_this_module(self, monkeypatch): + """The adapter method must not carry a second copy of the table. + + Checked by delegation rather than by constructing an ``ANI2xtAdapter``: + construction loads ``models/ani2xt_no_repulsion.pt`` and needs torchani + for the AEV computer, neither of which belongs in the fast tier. + """ + from Auto3D.models import adapter as adapter_mod + from Auto3D.models.adapter import ANI2xtAdapter + + seen: list = [] + + def _spy(atomic_numbers): + seen.append(list(atomic_numbers)) + return ["sentinel"] + + monkeypatch.setattr(adapter_mod, "to_ani2xt_species", _spy) + # Unbound call: no adapter instance, so no weights and no torchani. + result = ANI2xtAdapter.to_species(object(), [6, 1, 1, 1, 1]) + + assert seen == [[6, 1, 1, 1, 1]] + assert result == ["sentinel"] + + def test_the_mapping_itself_is_what_ani2xt_expects(self): + from Auto3D.models.species import to_ani2xt_species + + assert to_ani2xt_species([6, 1, 1, 1, 1]) == [1, 0, 0, 0, 0] + assert to_ani2xt_species([1, 6, 7, 8, 9, 16, 17]) == [0, 1, 2, 3, 4, 5, 6] + + def test_out_of_set_element_names_the_element_and_the_model(self): + from Auto3D.models.species import to_ani2xt_species + + with pytest.raises(ValueError) as exc: + to_ani2xt_species([11]) + message = str(exc.value) + assert "11" in message and "Na" in message and "ANI2xt" in message + + +def test_the_batch_opt_module_is_gone(): + """``batch_opt/species.py`` made the optimizer package a shared-utility host + for ``ASE/``, ``cli/`` and ``models/``'s own padder. Clean sweep, no alias.""" + import importlib + + with pytest.raises(ModuleNotFoundError): + importlib.import_module("Auto3D.batch_opt.species") diff --git a/tests/test_thermo_helpers.py b/tests/test_thermo_helpers.py index 9781ad58..442c0367 100644 --- a/tests/test_thermo_helpers.py +++ b/tests/test_thermo_helpers.py @@ -1265,7 +1265,12 @@ def _paramless_model(): import torch from torch import nn - class _ParamlessRecordingNNP(nn.Module): + from tests.helpers_adapter import AdapterModuleMixin + + # The mixin supplies the ModelAdapter members this double does not + # care about (pads, to_species, energy); EnForce_ANI gates on them. + # It contributes no nn.Parameter, so the premise below still holds. + class _ParamlessRecordingNNP(AdapterModuleMixin, nn.Module): def __init__(self): super().__init__() self.seen: list[dict] = [] diff --git a/tests/test_thermo_transition_state.py b/tests/test_thermo_transition_state.py index 9518d6bf..9ba3dd5c 100644 --- a/tests/test_thermo_transition_state.py +++ b/tests/test_thermo_transition_state.py @@ -149,7 +149,11 @@ def _zero_force_model(): import torch from torch import nn - class _StubNNP(nn.Module): + from tests.helpers_adapter import AdapterModuleMixin + + # The mixin supplies the ModelAdapter members this double does not + # care about (pads, to_species, energy); EnForce_ANI gates on them. + class _StubNNP(AdapterModuleMixin, nn.Module): def forward(self, coords, species, charges, atom_mask=None): energy = torch.zeros(coords.shape[0], dtype=coords.dtype) return energy, torch.zeros_like(coords).detach() diff --git a/tests/test_utils_validation.py b/tests/test_utils_validation.py index 75997c9e..35c7a7e2 100644 --- a/tests/test_utils_validation.py +++ b/tests/test_utils_validation.py @@ -256,119 +256,115 @@ def test_filter_unique_custom_threshold(self): assert len(result_strict) > len(result_lenient) +def _options(**overrides): + """Build the ``Auto3DOptions`` ``check_valid_configuration`` now takes. + + The function used to take ten keyword arguments mirroring + ``Auto3DOptions``'s field names *and carrying their own defaults* -- a third + configuration schema. It now takes the object, so these tests hand it one. + ``use_gpu=False`` by default: this box's CUDA availability must not decide + whether an assertion about paths or engines holds. + """ + from Auto3D.config import Auto3DOptions + + params = {"path": path_example_smi, "k": 1, "use_gpu": False} + params.update(overrides) + return Auto3DOptions(**params) + + class TestCheckValidConfiguration: """Tests for check_valid_configuration function.""" def test_valid_configuration(self): """Test that valid configuration returns no errors.""" errors = check_valid_configuration( - path=path_example_smi, - k=1, - use_gpu=False, - optimizing_engine="AIMNET", - isomer_engine="rdkit", - opt_steps=5000, + _options(optimizing_engine="AIMNET", isomer_engine="rdkit", opt_steps=5000) ) assert len(errors) == 0 def test_missing_path(self): """Test that missing path returns error.""" - errors = check_valid_configuration( - path=None, - k=1, - use_gpu=False, - ) + errors = check_valid_configuration(_options(path=None)) assert any("path" in e.lower() for e in errors) def test_nonexistent_path(self): """Test that nonexistent path returns error.""" - errors = check_valid_configuration( - path="/nonexistent/path.smi", - k=1, - use_gpu=False, - ) + errors = check_valid_configuration(_options(path="/nonexistent/path.smi")) assert any("exist" in e.lower() for e in errors) def test_missing_k_and_window(self): """Test that missing both k and window returns error.""" - errors = check_valid_configuration( - path=path_example_smi, - k=False, - window=False, - use_gpu=False, - ) + errors = check_valid_configuration(_options(k=False, window=False)) assert any("k" in e.lower() or "window" in e.lower() for e in errors) def test_window_specified(self): """Test that window alone is sufficient.""" - errors = check_valid_configuration( - path=path_example_smi, - k=False, - window=5.0, - use_gpu=False, - ) + errors = check_valid_configuration(_options(k=False, window=5.0)) assert not any("k" in e.lower() or "window" in e.lower() for e in errors) def test_invalid_optimizing_engine(self): - """Test that invalid optimizing_engine returns error.""" - errors = check_valid_configuration( - path=path_example_smi, - k=1, - use_gpu=False, - optimizing_engine="INVALID", - ) + """Test that invalid optimizing_engine returns error. + + Still checked here, unlike isomer_engine/tauto_engine below: an engine + name may be a registry entry or a path to a custom model, so it is not + an enumerable choice ``Auto3DOptions`` could validate from the value + alone -- it needs the registry lookup. + """ + errors = check_valid_configuration(_options(optimizing_engine="INVALID")) assert any("optimizing_engine" in e.lower() for e in errors) def test_accepts_aimnet_registry_names(self): """Registry engine names must validate, matching model_factory/CLI schema.""" for name in ("aimnet2", "aimnet2-2025", "aimnet2-nse", "aimnet2-pd"): - errors = check_valid_configuration( - path=path_example_smi, - k=1, - use_gpu=False, - optimizing_engine=name, - ) + errors = check_valid_configuration(_options(optimizing_engine=name)) assert not any("optimizing_engine" in e.lower() for e in errors), (name, errors) - def test_invalid_isomer_engine(self): - """Test that invalid isomer_engine returns error.""" - errors = check_valid_configuration( - path=path_example_smi, - k=1, - use_gpu=False, - isomer_engine="invalid", - ) - assert any("isomer_engine" in e.lower() for e in errors) + def test_invalid_isomer_engine_refused_at_construction(self): + """An unrecognized isomer_engine is refused before this function runs. - def test_opt_steps_too_small(self): - """Test that opt_steps < 10 returns error.""" - errors = check_valid_configuration( - path=path_example_smi, - k=1, - use_gpu=False, - opt_steps=5, - ) - assert any("opt_steps" in e.lower() for e in errors) + The whitelist moved to ``Auto3D.config.ENGINE_CHOICES`` and is enforced + by ``Auto3DOptions.__post_init__``, so ``check_valid_configuration`` can + no longer be reached with a bad value -- which is why it no longer + carries its own copy of the set. The rejection did not disappear; it + moved earlier, and to every entry point at once. + """ + from Auto3D.exceptions import ConfigurationError - def test_invalid_tauto_engine(self): - """Test that invalid tauto_engine with enumerate_tautomer=True returns error.""" - errors = check_valid_configuration( - path=path_example_smi, - k=1, - use_gpu=False, - enumerate_tautomer=True, - tauto_engine="invalid", - ) - assert any("tauto_engine" in e.lower() for e in errors) + with pytest.raises(ConfigurationError, match="isomer_engine"): + _options(isomer_engine="invalid") + + def test_opt_steps_too_small_refused_at_construction(self): + """opt_steps < 10 is refused at construction, not at run start. + + ``FIELD_BOUNDS["opt_steps"]`` is now ``("ge", 10)`` -- the single + declaration of that floor -- so the two hand-written ``< 10`` checks + that used to live in this module are gone. Same ``ConfigurationError``, + raised before the banner instead of after it. + """ + from Auto3D.exceptions import ConfigurationError + + with pytest.raises(ConfigurationError, match="opt_steps"): + _options(opt_steps=5) + + def test_invalid_tauto_engine_refused_at_construction(self): + """An unrecognized tauto_engine is refused at construction too. + + Unconditionally, where the old check only looked when + ``enumerate_tautomer`` was true -- ``CLIConfig``'s + ``Literal["rdkit", "oechem"]`` always rejected it, so the gated check + was an entry-point divergence. + """ + from Auto3D.exceptions import ConfigurationError + + with pytest.raises(ConfigurationError, match="tauto_engine"): + _options(enumerate_tautomer=True, tauto_engine="invalid") + with pytest.raises(ConfigurationError, match="tauto_engine"): + _options(enumerate_tautomer=False, tauto_engine="invalid") def test_valid_tauto_configuration(self): """Test valid tautomer configuration.""" errors = check_valid_configuration( - path=path_example_smi, - k=1, - use_gpu=False, - enumerate_tautomer=True, - tauto_engine="rdkit", + _options(enumerate_tautomer=True, tauto_engine="rdkit") ) # Should not have tauto_engine errors assert not any("tauto_engine" in e.lower() for e in errors) diff --git a/tests/test_validation.py b/tests/test_validation.py index ab90dfcc..f846d8a2 100644 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -5,6 +5,7 @@ import pytest from unittest.mock import patch, MagicMock +from Auto3D.config import Auto3DOptions from Auto3D.utils.validation import check_input from Auto3D.exceptions import GPUError, DependencyError, ConfigurationError, ModelLoadError @@ -104,17 +105,20 @@ def test_custom_nnp_load_failure_raises_model_load_error(self, tmp_path): check_input(args) def test_opt_steps_too_small_raises_configuration_error(self): - """Should raise ConfigurationError when opt_steps < 10.""" - args = MagicMock() - args.use_gpu = False - args.isomer_engine = "rdkit" - args.optimizing_engine = "AIMNET" - args.opt_steps = 5 - args.input_format = "smi" - args.path = "/fake/path.smi" - - with pytest.raises(ConfigurationError, match="smaller than 10"): - check_input(args) + """opt_steps < 10 must still raise ConfigurationError -- but now at + configuration *construction*, which is where the single declaration of + that floor lives. + + ``check_input`` used to hand-write ``< 10`` and so did + ``check_valid_configuration``, while ``Auto3D.config.FIELD_BOUNDS`` + declared ``("ge", 1)``: one option, two different minimums, and + ``opt_steps=5`` accepted by ``Auto3DOptions``/``CLIConfig`` only to fail + later. ``FIELD_BOUNDS["opt_steps"]`` is now ``("ge", 10)`` and both + hand-written checks are gone, so the refusal happens before any entry + point can act on the value. Same exception type, strictly earlier. + """ + with pytest.raises(ConfigurationError, match="opt_steps"): + Auto3DOptions(path="/fake/path.smi", k=1, opt_steps=5) def test_only_aimnet_molecules_with_ani2x_raises_configuration_error(self, tmp_path): """Should raise ConfigurationError when molecules require AIMNET but ANI2x selected.""" @@ -184,7 +188,9 @@ def test_out_of_range_index_flagged(self, tmp_path): p.write_text("CCO mol\n") with patch('Auto3D.utils.validation.torch.cuda.is_available', return_value=True), \ patch('Auto3D.utils.validation.torch.cuda.device_count', return_value=1): - errors = check_valid_configuration(path=str(p), k=1, use_gpu=True, gpu_idx=5) + errors = check_valid_configuration( + Auto3DOptions(path=str(p), k=1, use_gpu=True, gpu_idx=5) + ) assert any("GPU index 5 is invalid" in e for e in errors) def test_valid_index_not_flagged(self, tmp_path): @@ -193,7 +199,9 @@ def test_valid_index_not_flagged(self, tmp_path): p.write_text("CCO mol\n") with patch('Auto3D.utils.validation.torch.cuda.is_available', return_value=True), \ patch('Auto3D.utils.validation.torch.cuda.device_count', return_value=4): - errors = check_valid_configuration(path=str(p), k=1, use_gpu=True, gpu_idx=0) + errors = check_valid_configuration( + Auto3DOptions(path=str(p), k=1, use_gpu=True, gpu_idx=0) + ) assert errors == [] @@ -237,7 +245,7 @@ def test_check_valid_configuration_raises_gpu_error_not_configuration_error( p.write_text("CCO mol\n") with patch('Auto3D.utils.validation.torch.cuda.is_available', return_value=False): with pytest.raises(GPUError): - check_valid_configuration(path=str(p), k=1, use_gpu=True) + check_valid_configuration(Auto3DOptions(path=str(p), k=1, use_gpu=True)) def test_check_input_and_check_valid_configuration_agree(self, tmp_path): """The two entry points main() and smiles2mols reach check_input and @@ -259,7 +267,7 @@ def test_check_input_and_check_valid_configuration_agree(self, tmp_path): with pytest.raises(GPUError) as exc_via_check_input: check_input(args) with pytest.raises(GPUError) as exc_via_check_valid_configuration: - check_valid_configuration(path=str(p), k=1, use_gpu=True) + check_valid_configuration(Auto3DOptions(path=str(p), k=1, use_gpu=True)) assert type(exc_via_check_input.value) is type( exc_via_check_valid_configuration.value diff --git a/tests/test_workflow.py b/tests/test_workflow.py index 5fd996c2..1bed7b59 100644 --- a/tests/test_workflow.py +++ b/tests/test_workflow.py @@ -232,19 +232,11 @@ class TestOptimizerEmptyInput: def test_optimizer_handles_missing_file(self, tmp_path, caplog, monkeypatch): """Optimizer should gracefully handle missing input files.""" import logging - from types import SimpleNamespace import torch from Auto3D.batch_opt.batchopt import optimizing - - # This test exercises missing-file handling only -- run() returns before - # touching the model -- so stub create_model to skip the multi-second - # real AIMNet2 load (and stay robust to sibling tests clearing the cache). - monkeypatch.setattr( - "Auto3D.batch_opt.batchopt.create_model", - lambda *a, **k: SimpleNamespace(coord_pad=0.0, species_pad=-1), - ) + from tests.helpers_adapter import FakeAdapter device = torch.device("cpu") config = { @@ -255,7 +247,10 @@ def test_optimizer_handles_missing_file(self, tmp_path, caplog, monkeypatch): } nonexistent = str(tmp_path / "nonexistent.sdf") - optimizer = optimizing(nonexistent, str(tmp_path / "out.sdf"), "AIMNET", device, config) + # An injected double, because `optimizing` no longer constructs its own + # adapter -- and this test returns before the model is touched anyway. + optimizer = optimizing(nonexistent, str(tmp_path / "out.sdf"), + adapter=FakeAdapter(), device=device, config=config) # Should not raise, just log warning and return with caplog.at_level(logging.WARNING): @@ -266,18 +261,11 @@ def test_optimizer_handles_missing_file(self, tmp_path, caplog, monkeypatch): def test_optimizer_handles_empty_file(self, tmp_path, caplog, monkeypatch): """Optimizer should gracefully handle empty input files.""" import logging - from types import SimpleNamespace import torch from Auto3D.batch_opt.batchopt import optimizing - - # Empty-file handling returns before the model is used; stub create_model - # to skip the real AIMNet2 load (see test_optimizer_handles_missing_file). - monkeypatch.setattr( - "Auto3D.batch_opt.batchopt.create_model", - lambda *a, **k: SimpleNamespace(coord_pad=0.0, species_pad=-1), - ) + from tests.helpers_adapter import FakeAdapter device = torch.device("cpu") config = { @@ -291,7 +279,8 @@ def test_optimizer_handles_empty_file(self, tmp_path, caplog, monkeypatch): empty_sdf = tmp_path / "empty.sdf" empty_sdf.write_text("") - optimizer = optimizing(str(empty_sdf), str(tmp_path / "out.sdf"), "AIMNET", device, config) + optimizer = optimizing(str(empty_sdf), str(tmp_path / "out.sdf"), + adapter=FakeAdapter(), device=device, config=config) # Should not raise, just log warning and return with caplog.at_level(logging.WARNING): @@ -421,7 +410,7 @@ def test_optim_rank_wrapper_isolates_failing_chunks(tmp_path, monkeypatch): attempted = [] class _BoomOptimizing: - def __init__(self, in_f, out_f, engine, device, config, progress_cb=None): + def __init__(self, in_f, out_f, *, adapter, device, config, progress_cb=None): self._enumerated = in_f def run(self): @@ -978,18 +967,12 @@ def test_the_optimizer_names_each_record_it_could_not_parse( `SPE.calc_spe` and `ASE/thermo`'s `iter_thermo_records` both log per-record for exactly this; this reader did not. """ - from types import SimpleNamespace - import torch from rdkit import Chem from rdkit.Chem import AllChem from Auto3D.batch_opt.batchopt import optimizing - - monkeypatch.setattr( - "Auto3D.batch_opt.batchopt.create_model", - lambda *a, **k: SimpleNamespace(coord_pad=0.0, species_pad=-1), - ) + from tests.helpers_adapter import FakeAdapter mol = Chem.AddHs(Chem.MolFromSmiles("CCO")) AllChem.EmbedMolecule(mol, randomSeed=1) @@ -1006,8 +989,8 @@ def test_the_optimizer_names_each_record_it_could_not_parse( "batchsize_atoms": 1024, } optimizer = optimizing( - str(bad_sdf), str(tmp_path / "out.sdf"), "AIMNET", - torch.device("cpu"), config, + str(bad_sdf), str(tmp_path / "out.sdf"), adapter=FakeAdapter(), + device=torch.device("cpu"), config=config, ) with caplog.at_level(logging.WARNING):