From d250fe9cfcf2f6e2b3fd1183e7e7692cd787f9c7 Mon Sep 17 00:00:00 2001 From: isayev Date: Tue, 4 Aug 2026 00:56:08 -0400 Subject: [PATCH 1/6] test: give the isomers, utils, config and filtering tests real assertions These cover the code the next several refactors move, so they are hardened first: a test that asserts nothing cannot tell you whether a refactor preserved behavior, and hardening it afterwards proves nothing about the move. The substantive changes, rather than the mechanical ones: test_parallel_embed's dynamic-conformer test used ethane, which is degenerate and always collapses to a single conformer, so the assertion held for a reason unrelated to what it claimed. It now uses hexane and bounds the result against calculate_conformer_count. test_SDF2chunks compared only chunk counts; it now reparses each chunk and asserts the molecules match the source by name, order and atom count. test_config's chunk-meta test built its expectation by hand from the same literal it then compared against. It now derives from ChunkMeta's own annotations and required keys. Two tests are deleted rather than hardened. test_config_schema_exists cannot fail once the module imports, and no field on Auto3DOptions has a mutable default for test_immutable_default_list to guard. Two xfail(strict=True) tests pin defects found while deriving the correct assertions. Both fail today for the reason named in the marker: - filter_unique_optimized raises KeyError on a record with no E_tot property, while the legacy filter_unique tolerates it. The two filters were given a shared duplicate criterion but still disagree on malformed input, and the stricter one is the survivor. - opt_steps carries two minimums: FIELD_BOUNDS declares ("ge", 1) while validation.py hand-writes >= 10 in two places. The test asserts the two agree rather than picking a value, so whichever is authoritative can win. Verified: 1299 passed, 9 skipped, 2 xfailed. --- tests/test_cli_config_schema.py | 6 --- tests/test_cli_security.py | 11 ++++- tests/test_config.py | 75 +++++++++++++++++++++-------- tests/test_filtering.py | 74 +++++++++++++++++++++++++++- tests/test_isomer_engine.py | 28 +++++++++++ tests/test_isomers.py | 85 +++++++++++++++++++++++++++------ tests/test_parallel_embed.py | 26 ++++++++-- tests/test_utils_chemistry.py | 52 ++++++++++++++------ tests/test_utils_validation.py | 44 ++++++++++++----- 9 files changed, 326 insertions(+), 75 deletions(-) diff --git a/tests/test_cli_config_schema.py b/tests/test_cli_config_schema.py index 113ccd8e..047c5dfb 100644 --- a/tests/test_cli_config_schema.py +++ b/tests/test_cli_config_schema.py @@ -5,12 +5,6 @@ from pathlib import Path -def test_config_schema_exists(): - """Config schema class should exist.""" - from Auto3D.cli.config_schema import CLIConfig - assert CLIConfig is not None - - def test_config_defaults(): """Config should have sensible defaults.""" from Auto3D.cli.config_schema import CLIConfig diff --git a/tests/test_cli_security.py b/tests/test_cli_security.py index 7350eeb0..610519ef 100644 --- a/tests/test_cli_security.py +++ b/tests/test_cli_security.py @@ -107,15 +107,22 @@ def test_a_whitespace_only_file_is_refused(self, tmp_path): load_yaml_config(cfg) def test_a_top_level_list_is_refused(self, tmp_path): + """``match=`` anchors on the phrase unique to THIS guard's message + ("...its top level is a {type}.") -- "must contain a YAML mapping" is + also a substring of the empty-file message (config_schema.py:341), so + that alone cannot tell this test apart from the empty-file guard + firing by mistake (e.g. if the not-a-mapping check were deleted and + the empty-file check's message merely happened to also match). + """ cfg = _write(tmp_path, "- k\n- window\n") - with pytest.raises(ConfigurationError, match="must contain a YAML mapping"): + with pytest.raises(ConfigurationError, match="top level is a list"): load_yaml_config(cfg) def test_a_top_level_scalar_is_refused(self, tmp_path): cfg = _write(tmp_path, "just a bare string\n") - with pytest.raises(ConfigurationError, match="must contain a YAML mapping"): + with pytest.raises(ConfigurationError, match="top level is a str"): load_yaml_config(cfg) def test_unparseable_yaml_is_refused(self, tmp_path): diff --git a/tests/test_config.py b/tests/test_config.py index e8d1dc27..27376a07 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -125,34 +125,33 @@ def test_gpu_idx_list(self): config = Auto3DOptions(gpu_idx=[0, 1, 2]) assert config.gpu_idx == [0, 1, 2] - def test_immutable_default_list(self): - """Test that default list values are not shared between instances.""" - config1 = Auto3DOptions() - config2 = Auto3DOptions() - - # If gpu_idx is a list, modifying one shouldn't affect the other - # But with default int value, this just verifies independence - config1["k"] = 5 - assert config2.k is False # Default is False, not None - class TestChunkMeta: """Tests for ChunkMeta TypedDict.""" def test_chunk_meta_structure(self): - """Test that ChunkMeta can be used as expected.""" + """Track the TypedDict's OWN declared keys, not a hand-copied literal. + + The prior version built its own 5-key dict and asserted back the + values it had just set -- a key added to or removed from the real + ChunkMeta would never move this test. A TypedDict has no runtime + constructor to validate against, so pin it via ``__annotations__``/ + ``__required_keys__`` instead: a change to config.py's ChunkMeta now + has to be reflected here too, or this test fails. + """ from Auto3D.config import ChunkMeta - meta: ChunkMeta = { - "output": "/path/to/output.sdf", - "optimized_og": "/path/to/optimized.sdf", - "enumerated_sdf": "/path/to/enumerated.sdf", - "sorted_sdf": "/path/to/sorted.sdf", - "housekeeping_folder": "/path/to/housekeeping", + expected_keys = { + "output", "optimized_og", "output_taut", "smiles_enumerated", + "smiles_reduced", "smiles_hashed", "enumerated_sdf", "sorted_sdf", + "housekeeping_folder", "path", "dir", } - - assert meta["output"] == "/path/to/output.sdf" - assert meta["housekeeping_folder"] == "/path/to/housekeeping" + assert set(ChunkMeta.__annotations__) == expected_keys + # ChunkMeta declares no Optional/NotRequired fields, so every key is + # required -- a self-consistency check that would catch a field + # becoming optional without a matching intent. + assert ChunkMeta.__required_keys__ == frozenset(expected_keys) + assert ChunkMeta.__optional_keys__ == frozenset() def test_optimization_config_exposes_no_energy_criterion_knobs(): @@ -186,6 +185,42 @@ 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. + """ + from Auto3D.config import FIELD_BOUNDS, check_field_bounds + from Auto3D.utils.validation import check_valid_configuration + + kind, bound_min = FIELD_BOUNDS["opt_steps"] + assert kind == "ge" + + # config.py's own gate must accept its own declared floor (premise, not + # the point of this test). + 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}" + ) + + def test_negative_k_rejected(): from Auto3D.config import Auto3DOptions from Auto3D.exceptions import ConfigurationError diff --git a/tests/test_filtering.py b/tests/test_filtering.py index 354fe1dc..c3626466 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -61,10 +61,24 @@ def test_empty_list_returns_empty(self): assert result == [] def test_single_mol_returns_itself(self): - """Single molecule should be returned as-is.""" + """Single molecule should be returned as-is, explicit Hs and all. + + The RMSD comparison strips Hs from a throwaway copy for speed; the + returned molecule must be the caller's original (H-explicit) object, + not the no-H comparison copy -- the MLIP downstream requires explicit + H, and this is the len(mols) <= 1 short-circuit that never even + reaches the strip/compare loop. + """ mol = _create_mol_with_energy("C", -10.0) + n_atoms_before = mol.GetNumAtoms() + assert any(a.GetAtomicNum() == 1 for a in mol.GetAtoms()), "test premise: has explicit Hs" + result = _filter_within_cluster([mol], rmsd_threshold=0.5) + assert len(result) == 1 + assert result[0] is mol, "the single-mol short-circuit must return the original object" + assert result[0].GetNumAtoms() == n_atoms_before + assert any(a.GetAtomicNum() == 1 for a in result[0].GetAtoms()) def test_identical_mols_returns_one(self): """Identical molecules should be deduplicated to one.""" @@ -217,11 +231,17 @@ def test_empty_list_returns_empty(self): assert result == [] def test_filters_unconverged_structures(self): - """Unconverged structures should be filtered out.""" + """Unconverged structures should be filtered out -- and it must be + specifically the unconverged one that is gone, not just any one of + the two (e.g. a dedup bug that merged them for an unrelated reason + would also leave len(result) == 1). + """ mol1 = _create_mol_with_energy("C", -10.0, converged=True) mol2 = _create_mol_with_energy("CC", -11.0, converged=False) result = filter_unique_optimized([mol1, mol2], rmsd_threshold=0.5) assert len(result) == 1 + assert result[0] is mol1, "the converged structure must be the survivor" + assert result[0].GetProp("Converged").lower() == "true" def test_removes_duplicates(self): """Optimized filter should remove similar structures.""" @@ -311,6 +331,56 @@ def test_small_energy_window_creates_separate_clusters(self): assert len(result) == 2 +class TestMissingEnergyPropertyMustNotCrash: + """filter_unique_optimized must tolerate a record with no 'E_tot', the + way the legacy ``utils.chemistry.filter_unique`` already does. + + KNOWN DEFECT (found during cluster E brainstorming, not fixed by this + lane): ``filtering.py:75`` sorts the valid-mols list by + ``Auto3D.utils.energy.e_tot_ev``, which RAISES (KeyError/ValueError) for + a molecule with no usable 'E_tot' property. ``_filter_within_cluster``'s + own energy guard, two dozen lines later in the same module, instead uses + the tolerant ``try_e_tot_ev`` and treats a missing energy as "fall back + to RMSD only". ``utils.chemistry.filter_unique`` (the OTHER conformer + filter, sharing the same duplicate criterion since 4.0.1) also uses + ``try_e_tot_ev`` throughout and does not crash on this input. So the two + filters diverge on malformed input: the same list of mols that + ``filter_unique`` happily filters crashes ``filter_unique_optimized``. + + This matters here specifically because cluster B5 is about to delete one + of the two filters, and the survivor is the stricter (crashing) one -- + fixing filtering.py's sort key to use ``try_e_tot_ev``, matching its own + energy guard and the legacy filter, is what should make this pass. + """ + + @pytest.mark.xfail( + strict=True, + reason=( + "filtering.py:75 sorts by e_tot_ev (raises KeyError for a mol " + "with no 'E_tot' property) instead of the tolerant try_e_tot_ev " + "that _filter_within_cluster's own energy guard and the legacy " + "utils.chemistry.filter_unique both use -- the two conformer " + "filters disagree on malformed input (cluster E brainstorm defect)." + ), + ) + def test_missing_e_tot_property_does_not_crash(self): + mol_no_energy = Chem.AddHs(Chem.MolFromSmiles("CCO")) + AllChem.EmbedMolecule(mol_no_energy, randomSeed=42) + mol_no_energy.SetProp("Converged", "true") + # Deliberately no set_e_tot_from_ev call: this record has no 'E_tot'. + assert not mol_no_energy.HasProp("E_tot"), "test premise" + + mol_with_energy = _create_mol_with_energy("CC", -10.0) + + result = filter_unique_optimized( + [mol_no_energy, mol_with_energy], rmsd_threshold=0.3 + ) + + # Correct behavior: no crash, and a mol with no usable energy simply + # cannot be deduped by energy -- it must survive alongside the other. + assert len(result) == 2 + + class TestFilterUniqueBehavior: """Tests verifying behavior matches original filter_unique.""" diff --git a/tests/test_isomer_engine.py b/tests/test_isomer_engine.py index da92505c..393df1c1 100644 --- a/tests/test_isomer_engine.py +++ b/tests/test_isomer_engine.py @@ -92,9 +92,37 @@ def test_rd_isomer_conformer_func(): def test_SDF2chunks(): + """Chunks must partition the source records exactly: same count, same + identity (by name) in the same order, and same atom count each -- not + merely the same total count, which passes even if a chunk boundary + duplicated or dropped a molecule and another chunk absorbed the slack. + """ chunks = SDF2chunks(example_sdf) assert(len(chunks) == count_sdf(example_sdf)) + reference_mols = [ + mol for mol in Chem.SDMolSupplier(example_sdf, removeHs=False) + if mol is not None + ] + chunk_mols = [ + Chem.MolFromMolBlock("".join(chunk), removeHs=False) for chunk in chunks + ] + assert all(mol is not None for mol in chunk_mols), ( + "a chunk failed to parse back into a molecule" + ) + + reference_names = [m.GetProp("_Name") for m in reference_mols] + chunk_names = [m.GetProp("_Name") for m in chunk_mols] + assert chunk_names == reference_names, ( + "chunks do not reproduce the source records' identity/order" + ) + + reference_atom_counts = [m.GetNumAtoms() for m in reference_mols] + chunk_atom_counts = [m.GetNumAtoms() for m in chunk_mols] + assert chunk_atom_counts == reference_atom_counts, ( + "a chunk's atom count does not match its source record" + ) + def test_rd_isomer_with_parallel_embedding(): """Test RDKitIsomer with parallel embedding enabled.""" diff --git a/tests/test_isomers.py b/tests/test_isomers.py index dc8e7905..8a23fda1 100644 --- a/tests/test_isomers.py +++ b/tests/test_isomers.py @@ -170,15 +170,23 @@ def test_unknown_engine_raises_error(self): ) def test_engine_type_case_insensitive(self): - """Test that engine type is case insensitive.""" - # This should not raise - it will fail on actual instantiation - # but the case normalization should work - with pytest.raises(ValueError, match="Unknown isomer engine type"): - create_isomer_engine( - "UNKNOWN", + """A *valid* engine name in an unexpected case must resolve to the + correct adapter, not merely fail to crash on an already-invalid name. + + The previous version passed "UNKNOWN" -- invalid in any case -- so it + could never have distinguished case normalization working from case + normalization being entirely absent. + """ + for name in ("RDKit", "RDKIT", "rdkit"): + engine = create_isomer_engine( + name, input_path="/input.smi", output_path="/output.sdf", + smiles_enumerated="/enum.smi", + smiles_reduced="/reduced.smi", + smiles_hashed="/hashed.smi", ) + assert isinstance(engine, RDKitIsomerAdapter), name def test_omega_engine_creates_adapter(self): """Test that 'omega' creates OmegaIsomerAdapter.""" @@ -257,6 +265,52 @@ def test_rdkit_engine_parallel_embedding_enabled(self, tmp_path): assert engine.parallel_embedding_threshold == 5 assert engine.parallel_workers == 2 + def test_rdkit_engine_parallel_embedding_enabled_actually_runs_parallel_path( + self, tmp_path, monkeypatch + ): + """Constructor kwargs alone don't prove the parallel path executes. + + Drive ``.run()`` for real (small, hermetic, no NNP) and spy on + ``embed_conformers_parallel`` -- the parallel path's only entry point + -- so a regression that silently falls back to serial embedding + would be caught even though every attribute above still reports + correctly. + """ + import Auto3D.isomers.parallel_embed as parallel_embed_mod + + job_dir = tmp_path / "job" + job_dir.mkdir() + smi = tmp_path / "in.smi" + smi.write_text("CCO ethanol\n") + + calls = {"n": 0} + + def spy(*args, **kwargs): + calls["n"] += 1 + return iter([]) # no conformers written; only the call matters + + monkeypatch.setattr(parallel_embed_mod, "embed_conformers_parallel", spy) + + engine = create_isomer_engine( + "rdkit", + input_path=str(smi), + output_path=str(tmp_path / "output.sdf"), + smiles_enumerated=str(tmp_path / "enum.smi"), + smiles_reduced=str(tmp_path / "reduced.smi"), + smiles_hashed=str(tmp_path / "hashed.smi"), + job_dir=str(job_dir), + use_parallel_embedding=True, + parallel_embedding_threshold=1, # even one molecule triggers it + parallel_workers=2, + ) + + engine.run() + + assert calls["n"] == 1, ( + "embed_conformers_parallel was never called: the parallel path " + "did not run despite use_parallel_embedding=True" + ) + class TestCreateTautomerEngine: """Tests for create_tautomer_engine factory function.""" @@ -271,14 +325,17 @@ def test_unknown_engine_raises_error(self): ) def test_engine_type_case_insensitive(self): - """Test that engine type is case insensitive for valid types.""" - # Both RDKIT and rdkit should work (case normalization happens) - with pytest.raises(ValueError, match="Unknown tautomer engine type"): - create_tautomer_engine( - "UNKNOWN", - input_path="/input.smi", - output_path="/output.smi", - ) + """A *valid* engine name ("RDKIT") must resolve to the same + rdkit-backed engine as lowercase "rdkit" -- not merely fail to crash + on an already-invalid name, which "UNKNOWN" could never distinguish. + """ + from Auto3D.isomer_engine import TautomerEngine as TautEngine + + engine = create_tautomer_engine( + "RDKIT", input_path="/input.smi", output_path="/output.smi" + ) + assert isinstance(engine, TautEngine) + assert engine.mode == "rdkit" class TestIsomerEngineFactory: diff --git a/tests/test_parallel_embed.py b/tests/test_parallel_embed.py index 7643f407..40b5e1a6 100644 --- a/tests/test_parallel_embed.py +++ b/tests/test_parallel_embed.py @@ -39,16 +39,34 @@ def test_embed_single_returns_list_of_tuples(self): assert "methane" in conf_id def test_embed_single_with_dynamic_conformers(self): - """_embed_single with n_conformers=None should use dynamic calculation.""" + """_embed_single with n_conformers=None should use dynamic calculation. + + Compare the actual conformer count against ``calculate_conformer_count``'s + own formula for this molecule, instead of a bare ">= 1" that would + pass even if the None branch silently stopped calling that formula + (e.g. fell back to a fixed conformer count). Hexane is flexible + enough that the two counts are not degenerate (unlike a rigid/small + molecule, where embedding + RMSD pruning collapses to 1 regardless of + the requested count and could hide a formula regression). + """ + from Auto3D.utils.chemistry import calculate_conformer_count + + mol = Chem.AddHs(Chem.MolFromSmiles("CCCCCC")) # hexane: flexible + expected_upper_bound = calculate_conformer_count(mol) + results = _embed_single( - smi="CC", - name="ethane", + smi="CCCCCC", + name="hexane", n_conformers=None, threshold=0.3, np_threads=1, ) - assert len(results) >= 1 + assert 1 <= len(results) <= expected_upper_bound, ( + f"expected between 1 and {expected_upper_bound} conformers " + "(calculate_conformer_count's own dynamic formula for hexane), " + f"got {len(results)}" + ) def test_embed_single_filters_invalid_conformers(self): """_embed_single should filter conformers with atom clashes.""" diff --git a/tests/test_utils_chemistry.py b/tests/test_utils_chemistry.py index e5b3c60c..f0b8345a 100644 --- a/tests/test_utils_chemistry.py +++ b/tests/test_utils_chemistry.py @@ -370,27 +370,44 @@ def test_amend_mol_returns_none_for_invalid(self): # (either returns None or attempts to fix) def test_amend_mol_with_sanitize(self): - """Test that amend_mol can sanitize molecules.""" + """amend_mol(sanitize=True) must actually run RDKit's sanitization, + not just return a non-None object. + + A molecule parsed with ``sanitize=False`` has no ring-perception / + implicit-valence pass yet, so querying ring info raises a + precondition violation -- it genuinely needs sanitizing to become + usable. If ``amend_mol``'s sanitize branch were a no-op, this + molecule would still raise after the call. + """ from Auto3D.utils.chemistry import amend_mol - mol = Chem.MolFromSmiles("CCO") - mol = Chem.AddHs(mol) - AllChem.EmbedMolecule(mol, randomSeed=42) + + mol = Chem.MolFromSmiles("c1ccccc1", sanitize=False) # benzene, unsanitized + with pytest.raises(RuntimeError): + mol.GetRingInfo().NumRings() # ring perception never ran amended_mol = amend_mol(mol, sanitize=True) + assert amended_mol is not None + # Sanitizing actually ran: ring perception now works and finds the ring. + assert amended_mol.GetRingInfo().NumRings() == 1 class TestGetMolConnectivity: """Test the get_mol_connectivity function.""" def test_ethane_connectivity(self): - """Test connectivity for ethane (C-C single bond).""" + """Test connectivity for ethane (C-C single bond). + + Pins the exact canonical ordering (atom1_idx < atom2_idx, per the + function's own docstring/example), not "either order" -- which would + equally accept a broken ``get_mol_connectivity`` that stopped sorting + its tuples. + """ from Auto3D.utils.chemistry import get_mol_connectivity mol = Chem.MolFromSmiles("CC") connectivity = get_mol_connectivity(mol) - # Should have C-C bond - assert (0, 1) in connectivity or (1, 0) in connectivity + assert connectivity == {(0, 1)} def test_ethanol_connectivity(self): """Test connectivity for ethanol.""" @@ -425,18 +442,22 @@ def test_methane_connectivity(self): assert len(connectivity_with_h) == 4 # 4 C-H bonds def test_include_bond_order(self): - """Test that bond order can be included.""" + """Test that bond order can be included. + + The previous version's real assertion sat inside ``if len(bond_info) + == 3:``, which is false exactly when ``include_bond_order`` silently + stops adding the third element -- the one failure mode this test + exists to catch. Assert the 3-tuple shape unconditionally, then the + bond order value. + """ from Auto3D.utils.chemistry import get_mol_connectivity mol = Chem.MolFromSmiles("C=C") # Ethene connectivity = get_mol_connectivity(mol, include_bond_order=True) - # Should contain tuple with bond order info - # Format: (atom1_idx, atom2_idx, bond_order) - assert len(connectivity) >= 1 + assert connectivity == {(0, 1, 2.0)} for bond_info in connectivity: - if len(bond_info) == 3: - # Has bond order - assert bond_info[2] == 2.0 # Double bond + assert len(bond_info) == 3 # (atom1_idx, atom2_idx, bond_order) + assert bond_info[2] == 2.0 # Double bond class TestFilterUnique: @@ -633,6 +654,9 @@ def test_filter_custom_threshold(self): # Large threshold should definitely merge identical mols assert len(unique_mols_large) == 1 + # A tighter threshold can never merge MORE than a looser one -- the + # discarded half of this test's own computation, now actually checked. + assert len(unique_mols_small) >= len(unique_mols_large) def test_filter_unique_removehs_is_linear_and_nondestructive(self, monkeypatch): """Legacy filter_unique strips Hs once per molecule (not per comparison) and diff --git a/tests/test_utils_validation.py b/tests/test_utils_validation.py index 732380f3..75997c9e 100644 --- a/tests/test_utils_validation.py +++ b/tests/test_utils_validation.py @@ -223,19 +223,37 @@ def test_filter_unique_keeps_records_with_no_converged_property(self): ) def test_filter_unique_custom_threshold(self): - """Test filter_unique with custom RMSD threshold.""" - supp = Chem.SDMolSupplier(path_example_sdf, removeHs=False) - mols = [mol for mol in supp if mol is not None] - - for mol in mols: - mol.SetProp("Converged", "True") - - # With very small threshold, more structures should be kept - result_strict = filter_unique(mols, crit=0.01) - # With larger threshold, fewer structures should be kept - result_lenient = filter_unique(mols, crit=1.0) - - assert len(result_strict) >= len(result_lenient) + """A tighter RMSD threshold must keep MORE structures than a looser + one -- not merely "at least as many", which passes even when the two + thresholds produce identical results (as the fixture in + ``path_example_sdf`` does: two molecules of different sizes, so + ``species_key`` alone already keeps both regardless of ``crit``). + + Constructs two conformers of ONE molecule whose RMSD sits strictly + between the two thresholds, so equality cannot pass silently. + """ + mol1 = Chem.AddHs(Chem.MolFromSmiles("CCCCCCCC")) # octane: flexible + mol2 = Chem.Mol(mol1) + from rdkit.Chem import AllChem, rdMolAlign + AllChem.EmbedMolecule(mol1, randomSeed=1) + AllChem.EmbedMolecule(mol2, randomSeed=99) + AllChem.MMFFOptimizeMolecule(mol1) + AllChem.MMFFOptimizeMolecule(mol2) + mol1.SetProp("Converged", "True") + mol2.SetProp("Converged", "True") + + rmsd = rdMolAlign.GetBestRMS(Chem.RemoveHs(mol1), Chem.RemoveHs(mol2)) + assert rmsd > 0.05, "test premise: conformers must be geometrically distinct" + + crit_strict = rmsd / 2 # below the actual RMSD -> kept separate + crit_lenient = rmsd * 2 # above the actual RMSD -> merged + + result_strict = filter_unique([mol1, mol2], crit=crit_strict) + result_lenient = filter_unique([mol1, mol2], crit=crit_lenient) + + assert len(result_strict) == 2, "strict threshold must not merge distinct conformers" + assert len(result_lenient) == 1, "lenient threshold must merge near-identical conformers" + assert len(result_strict) > len(result_lenient) class TestCheckValidConfiguration: From e31b5a829edea12ac60428508c8609e9ee16e9f0 Mon Sep 17 00:00:00 2001 From: isayev Date: Tue, 4 Aug 2026 00:56:52 -0400 Subject: [PATCH 2/6] perf!: stop importing torch and rdkit to import Auto3D Three eager optional-dependency probes in __init__.py defeated the _LAZY_API mechanism that exists to avoid exactly this. Probing for ANI2xt reached batch_opt, which reached the utils barrel, which reached validation, which imported torch and models.loading. So `import Auto3D` paid for the entire package plus torch and rdkit before the caller had asked for anything. Measured in a subprocess on one box: import Auto3D 1.35 s -> 0.031 s len(sys.modules) 1175 -> 154 (stdlib floor is 128) torch, rdkit loaded yes -> no Auto3D.* submodules 20 -> 0 The probes are deleted rather than made lazy: they are not public API, and every real probe already exists at its use site. A comment records that and says not to reintroduce one. __version__ now comes from a private _detect_version(), and __dir__ reports the lazy API, so `main` is discoverable and warnings, version, ANI2xt and PackageNotFoundError stop leaking into the namespace. __getattr__ deliberately does not cache. Caching would turn it into a capturing binding, which is how a lazily imported reader ends up holding a stub permanently while monkeypatch reports success; test_lazy_torchani_import documents that failure appearing 182 tests downstream. A test asserts the non-caching property. The cost tests assert module count rather than wall-clock seconds, and run in a subprocess since the parent pytest process has already imported everything. Verified: five tripwires confirmed red first for the stated reasons, then each fix mutation-verified by reverting it and confirming the named test goes red. 1298 passed across three random seeds plus fixed order; ruff clean. --- src/Auto3D/__init__.py | 84 ++++++++--- tests/test_import_boundaries.py | 252 ++++++++++++++++++++++++++++++++ 2 files changed, 315 insertions(+), 21 deletions(-) create mode 100644 tests/test_import_boundaries.py diff --git a/src/Auto3D/__init__.py b/src/Auto3D/__init__.py index 54501df8..613017a6 100644 --- a/src/Auto3D/__init__.py +++ b/src/Auto3D/__init__.py @@ -5,32 +5,45 @@ using neural network potentials (AIMNet2, ANI2x, ANI2xt). """ -import warnings -from importlib.metadata import PackageNotFoundError, version -try: - __version__ = version(__name__) -except PackageNotFoundError: - __version__ = "unknown" +def _detect_version() -> str: + """Read the installed distribution version, or "unknown" if not installed. -# Optional dependency imports with proper exception handling -with warnings.catch_warnings(): - warnings.simplefilter("ignore") + The ``importlib.metadata`` import lives inside this function on purpose: at + module level it would put ``version`` and ``PackageNotFoundError`` into the + package namespace, where they are reachable as ``Auto3D.version`` and + ``Auto3D.PackageNotFoundError`` -- two names this package never meant to + export. See ``tests/test_import_boundaries.py``. + """ + from importlib.metadata import PackageNotFoundError, version try: - from openeye import oechem, oeomega, oequacpac # noqa: F401 (optional dependency probe) - except ImportError: - pass # OpenEye is optional + return version(__name__) + except PackageNotFoundError: + return "unknown" - try: - import torchani # noqa: F401 (optional dependency probe) - except ImportError: - pass # TorchANI is optional - try: - from Auto3D.batch_opt.ANI2xt_no_rep import ANI2xt # noqa: F401 (optional dependency probe) - except ImportError: - pass # ANI2xt model is optional +__version__ = _detect_version() + +# NOTE: this module imports nothing but the standard library, and that is a +# tested property, not an accident (tests/test_import_boundaries.py). Importing +# the package root is the cost every consumer pays unconditionally -- `auto3d +# --help`, a build script reading `__version__`, a tool that merely lists +# installed distributions -- and none of them need torch or rdkit. +# +# This file used to end with three eager optional-dependency probes (openeye, +# torchani, and `from Auto3D.batch_opt.ANI2xt_no_rep import ANI2xt`) wrapped in +# `warnings.catch_warnings()`. Nothing consumed any of them, and the third one +# defeated the `_LAZY_API` mechanism below outright: it reached ANI2xt_no_rep -> +# the Auto3D.utils barrel -> utils.validation -> Auto3D.models.* + torch, which +# turned `import Auto3D` into 1175 modules and 1.35 s and eagerly loaded 20 +# Auto3D submodules. Every probe is already duplicated where it is load-bearing +# and where its result is actually used: +# * openeye -> isomer_engine.py (names used at call time) and +# utils/validation.py (raises DependencyError with a fix hint) +# * torchani -> utils/validation.py (same) and batch_opt/ANI2xt_no_rep.py +# * ANI2xt -> constructed only through model_factory / models.adapter +# Do not reintroduce a probe here. A dependency is checked where it is needed. __all__ = [ "__version__", @@ -73,9 +86,38 @@ # Lazy imports for public API def __getattr__(name: str): - """Lazy import for public API functions (see _LAZY_API).""" + """Lazy import for public API functions (see _LAZY_API). + + Design constraint, not a missed optimization: this **must not** cache the + resolved object into ``globals()``. Caching would turn every access after + the first into a snapshot, i.e. exactly the import-time binding that makes + ``from X import y`` capture a stub -- a test that touches ``Auto3D.main`` + and then patches ``Auto3D.auto3D.main`` would patch nothing, report + success, and fail somewhere else entirely (the mechanism written up at + length in ``tests/test_lazy_torchani_import.py``). After the first access + ``import_module`` is a ``sys.modules`` dict lookup, so a cache buys nothing + measurable. Pinned by + ``test_import_boundaries.py::test_getattr_does_not_cache_resolved_attributes``. + """ if name in _LAZY_API: import importlib module_name, attr = _LAZY_API[name] return getattr(importlib.import_module(module_name), attr) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + """Report the public API alongside whatever is genuinely present. + + PEP 562 requires this to accompany ``__getattr__``: a module-level + ``__getattr__`` resolves names that are not in ``globals()``, and the + default ``dir()`` only sees ``globals()`` -- so without this, ``"main" in + dir(Auto3D)`` was False and neither tab-completion nor introspection could + find the public API. + + The union (rather than ``sorted(__all__)``) matters because ``__dir__`` + replaces the default entirely: imported submodules become real attributes + of the package, so ``Auto3D.cli`` after ``import Auto3D.cli`` must stay + visible. + """ + return sorted(set(globals()) | set(__all__)) diff --git a/tests/test_import_boundaries.py b/tests/test_import_boundaries.py new file mode 100644 index 00000000..d70e35ce --- /dev/null +++ b/tests/test_import_boundaries.py @@ -0,0 +1,252 @@ +"""Import boundaries for the ``Auto3D`` package root: what ``import Auto3D`` +is allowed to cost, and what the package root is allowed to expose. + +Two properties are locked here. + +**Import cost.** ``import Auto3D`` must reach nothing but the standard library. +The package root is the entry point every consumer pays for -- ``auto3d +--help``, ``from Auto3D import __version__`` in a build script, an unrelated +library that merely lists installed packages -- and none of those need torch, +rdkit, or a neural network potential. The cost is asserted as a **module count** +rather than a wall-clock time because seconds are a property of the machine +(CPU, warm page cache, NFS-mounted site-packages) while ``len(sys.modules)`` is +a property of the code. Eager optional-dependency probes in ``__init__.py`` once +made this 1175 modules / 1.35 s; the cap below leaves generous headroom over the +stdlib-only floor so it tracks a regression in kind, not a fluctuation in +degree. + +**Public surface.** ``__getattr__`` without ``__dir__`` violates PEP 562: +``dir()`` stops reporting the lazily resolved names (``"main" in dir(Auto3D)`` +was False) while the module namespace still exposes whatever the module body +happened to import. Both halves are asserted -- the public names are visible, +and the import-machinery names are not. + +Every cost measurement runs in a **subprocess**. It cannot be done in-process: +``conftest.py`` deliberately imports every ``Auto3D`` submodule before the first +test runs (see ``_import_every_auto3d_module_before_any_test``), so by the time +any assertion here executes, ``sys.modules`` already holds torch, rdkit, and all +of Auto3D. An in-process version of these tests would pass unconditionally. +""" +from __future__ import annotations + +import importlib +import json +import subprocess +import sys + +import pytest + +# Generous cap over the stdlib-only floor (~140 modules: interpreter startup +# plus importlib.metadata's own imports). Set well above the floor and well +# below the pre-fix 1175 so it survives a stdlib growing an internal import or +# a site-packages .pth adding one, while still going red the moment +# ``Auto3D/__init__.py`` reaches for torch, rdkit, or a domain submodule. +MAX_MODULES_AFTER_BARE_IMPORT = 250 + +# Names that leaked out of the package root as attributes without ever being +# part of the public API: the ``import warnings`` and ``from +# importlib.metadata import PackageNotFoundError, version`` used to compute +# ``__version__``, and the optional-dependency probes' bindings. +LEAK_CANDIDATES = ( + "warnings", + "version", + "PackageNotFoundError", + "ANI2xt", + "oechem", + "oeomega", + "oequacpac", + "torchani", +) + +# Runs in a fresh interpreter; reports what a bare ``import Auto3D`` cost and +# what it left reachable. Emits one JSON line on stdout. +_PROBE_SOURCE = """ +import json, sys, time + +t0 = time.perf_counter() +import Auto3D +elapsed = time.perf_counter() - t0 + +leak_candidates = %(leaks)r +print(json.dumps({ + "elapsed": elapsed, + "n_modules": len(sys.modules), + "auto3d_modules": sorted( + m for m in sys.modules if m == "Auto3D" or m.startswith("Auto3D.") + ), + "torch": any(m == "torch" or m.startswith("torch.") for m in sys.modules), + "rdkit": any(m == "rdkit" or m.startswith("rdkit.") for m in sys.modules), + "version": Auto3D.__version__, + "dir": sorted(dir(Auto3D)), + "all": list(Auto3D.__all__), + "reachable_leaks": sorted(n for n in leak_candidates if hasattr(Auto3D, n)), +})) +""" % {"leaks": LEAK_CANDIDATES} + + +@pytest.fixture(scope="module") +def bare_import(): + """Measure a cold ``import Auto3D`` in a fresh interpreter, once.""" + proc = subprocess.run( + [sys.executable, "-c", _PROBE_SOURCE], + capture_output=True, + text=True, + ) + assert proc.returncode == 0, f"probe failed:\n{proc.stdout}\n{proc.stderr}" + return json.loads(proc.stdout.strip().splitlines()[-1]) + + +# --------------------------------------------------------------------------- # +# Import cost +# --------------------------------------------------------------------------- # + +def test_bare_import_does_not_load_torch_or_rdkit(bare_import): + """The two heavyweight third-party dependencies stay unimported. + + They are what makes the difference between 0.02 s and 1.35 s. Anything that + needs them imports them itself, at the point of need. + """ + assert not bare_import["torch"], "import Auto3D pulled in torch" + assert not bare_import["rdkit"], "import Auto3D pulled in rdkit" + + +def test_bare_import_loads_no_auto3d_submodule(bare_import): + """Only the package root itself is imported. + + This is the sharp form of the cost assertion: it names the defect (the + package root executing a domain submodule's module body) rather than a + number. ``_LAZY_API`` exists precisely so that no submodule loads until a + public name is used. + """ + assert bare_import["auto3d_modules"] == ["Auto3D"], ( + "import Auto3D eagerly imported submodules: " + f"{[m for m in bare_import['auto3d_modules'] if m != 'Auto3D']}" + ) + + +def test_bare_import_module_count_stays_under_cap(bare_import): + """``len(sys.modules)`` after a bare import stays near the stdlib floor.""" + assert bare_import["n_modules"] < MAX_MODULES_AFTER_BARE_IMPORT, ( + f"import Auto3D loaded {bare_import['n_modules']} modules " + f"(cap {MAX_MODULES_AFTER_BARE_IMPORT}, took " + f"{bare_import['elapsed']:.3f}s)" + ) + + +def test_bare_import_still_reports_a_version(bare_import): + """Cutting the cost must not cost ``__version__``. + + ``__version__`` is computed by a private helper so that ``version`` and + ``PackageNotFoundError`` do not land in the module namespace; this checks + the helper still runs, in a fresh interpreter, outside pytest. + """ + assert isinstance(bare_import["version"], str) + assert bare_import["version"] + + +# --------------------------------------------------------------------------- # +# Public surface: __dir__ (PEP 562) +# --------------------------------------------------------------------------- # + +def test_dir_reports_the_public_api(): + """Every ``__all__`` name is visible to ``dir()`` and to tab-completion.""" + import Auto3D + + listed = set(dir(Auto3D)) + assert "main" in listed + missing = sorted(set(Auto3D.__all__) - listed) + assert not missing, f"__all__ names absent from dir(Auto3D): {missing}" + + +def test_dir_does_not_leak_import_machinery(bare_import): + """The package root exposes no name it merely needed in order to load. + + Checked in both directions and in both processes: not in ``dir()`` here, + and not reachable via ``hasattr`` in a fresh interpreter. ``dir()`` alone is + not enough -- a name can be absent from ``dir()`` and still resolve -- and + ``hasattr`` alone is not enough in this process, because ``conftest``'s + eager submodule import legitimately adds submodule attributes. + """ + import Auto3D + + listed = set(dir(Auto3D)) + leaked_here = sorted(n for n in LEAK_CANDIDATES if n in listed) + assert not leaked_here, f"dir(Auto3D) leaks import machinery: {leaked_here}" + assert not bare_import["reachable_leaks"], ( + "package root exposes non-public attributes: " + f"{bare_import['reachable_leaks']}" + ) + + +# --------------------------------------------------------------------------- # +# Public surface: _LAZY_API +# --------------------------------------------------------------------------- # + +def test_lazy_api_is_a_bijection_with_all(): + """``_LAZY_API`` and ``__all__`` describe the same surface. + + Without this, a name can be added to ``__all__`` and be unreachable, or + added to ``_LAZY_API`` and be undocumented. + """ + import Auto3D + + assert set(Auto3D._LAZY_API) == set(Auto3D.__all__) - {"__version__"} + + +def test_every_lazy_api_name_resolves_to_its_target(): + """Each lazy name resolves to the exact object at its declared location. + + Stronger than "is not None": it pins the target, so moving a function + without updating ``_LAZY_API`` fails here instead of silently exporting + something else of the same name. + """ + import Auto3D + + for name, (module_name, attr) in Auto3D._LAZY_API.items(): + target = getattr(importlib.import_module(module_name), attr) + assert getattr(Auto3D, name) is target, name + + +def test_getattr_does_not_cache_resolved_attributes(): + """Repeated access must re-read the source module, not return a snapshot. + + Caching into ``globals()`` looks like a free optimization and is not: it + turns ``__getattr__`` into an import-time binding, so a test that touches + ``Auto3D.main`` and *then* patches ``Auto3D.auto3D.main`` would patch + nothing and pass for the wrong reason -- the failure mode + ``tests/test_lazy_torchani_import.py`` documents, which surfaced 182 tests + downstream of its cause. After the first access ``import_module`` is a + ``sys.modules`` dict hit, so there is nothing to buy. + + Verified behaviorally (a post-access change to the source module is seen) + and structurally (the name never lands in the package namespace). + """ + import Auto3D + + first = Auto3D.Auto3DOptions + assert first is not None + assert "Auto3DOptions" not in vars(Auto3D), ( + "__getattr__ cached into the module namespace" + ) + + sentinel = object() + module = importlib.import_module("Auto3D.config") + original = module.Auto3DOptions + try: + module.Auto3DOptions = sentinel + assert Auto3D.Auto3DOptions is sentinel, ( + "second access returned a cached value instead of re-reading " + "Auto3D.config" + ) + finally: + module.Auto3DOptions = original + + assert Auto3D.Auto3DOptions is original + + +def test_unknown_attribute_raises_attribute_error(): + """``__getattr__`` must not turn a typo into an ImportError or a hang.""" + import Auto3D + + with pytest.raises(AttributeError, match="no attribute 'nope'"): + getattr(Auto3D, "nope") # noqa: B009 (the lookup is the assertion) From daa7952bd4f31a1ba1c8f2af7d4aec7b16ca1076 Mon Sep 17 00:00:00 2001 From: isayev Date: Tue, 4 Aug 2026 00:57:16 -0400 Subject: [PATCH 3/6] docs: record the import-cost change and the four removed namespace attributes Auto3D.ANI2xt, .warnings, .version and .PackageNotFoundError were never public API, but removing them is still visible to anyone who reached for them, so they belong in Breaking Changes rather than only in the perf note. --- CHANGELOG.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5eb8c64c..c648faae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 mass of the isotope it names; the change applies only where no isotope was specified. +- **`import Auto3D` no longer imports torch or RDKit, and four attributes are + gone from the package namespace.** `Auto3D.ANI2xt`, `Auto3D.warnings`, + `Auto3D.version` and `Auto3D.PackageNotFoundError` were never public API — + they were leaked into the namespace by three eager optional-dependency probes + and by module-level imports. All four are removed. + + The probes existed to detect whether an optional engine was installed, but + probing for ANI2xt reached `batch_opt`, which reached the `utils` barrel, + which reached `validation`, which imported torch and `models.loading`. So + importing the package paid for the whole package, plus torch and RDKit, + before the caller had asked for anything. Every real probe already exists at + its use site, so nothing was lost by deleting them. + + | | before | after | + |---|---|---| + | `import Auto3D` | 1.35 s | 0.031 s | + | `len(sys.modules)` | 1175 | 154 | + | torch / RDKit loaded | yes | no | + | `Auto3D.*` submodules loaded | 20 | 0 | + + **What stops working:** referencing any of those four names through the + `Auto3D` package object. **What to do instead:** `import warnings` yourself; + read the version from `Auto3D.__version__`; import optional engines from the + module that owns them. Documented public names are unaffected and still + resolve lazily on first access. + - **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 From 5dcc965475fbc959402db671402b0b94ac8f3291 Mon Sep 17 00:00:00 2001 From: isayev Date: Tue, 4 Aug 2026 01:04:39 -0400 Subject: [PATCH 4/6] test: make the model, batch_opt and pipeline tests able to fail The same hardening as the previous commit, for the tests covering the code the model-contract consolidation is about to move. The assertions that were doing no work: - test_create_unknown_model_raises_error caught bare Exception, so it could not tell an unresolvable registry name from any other failure. It now matches the message. - Two pytest.raises(ModelLoadError) had no match=, so either would have passed on an unrelated ModelLoadError. Both now match on the argument order they exist to check. - test_ensemble_opt_returns_convergence_info ignored the convergence info it is named for; it now asserts the mask and the oscillation counts. - test_fire_independent_molecule_tracking asserted that per-molecule state existed, not that it was independent. Rewritten as a two-phase protocol: bootstrap with every molecule progressing, then make exactly one oscillate, and assert the exact a and dt values that follow. - test_optimizer_handles_empty_file accepted any log output; it now asserts the empty-file message and that the missing-file message is absent, since those two guards were indistinguishable. New coverage where a code path had none: - The arity check's len(required) > 3 half was never exercised. Only the too-few-arguments half was. - The alias vocabulary is now parametrized over the full synonym set in both correct and transposed order, 14 cases, calling validate_custom_nnp directly so nothing needs to pickle. - forces = -grad is written separately in ANI2xtAdapter and ANI2xAdapter and was tested in neither. Both now assert -2*coords against a toy quadratic model, hermetically, by bypassing each adapter's real __init__. - The "produced no conformers after clash relief" warning had no test. One drives the real embedding path with relieve_clash stubbed by atom count, and asserts the warning fires once, names the right molecule, and that a sibling species survives without being warned about. - test_one_bad_molecule_does_not_remove_the_others parsed IDs with a stale split("_")[0] that made a Tier 1 test vacuous. test_calc_spe_uses_model_factory and test_model_name2model_calculator_uses_factory were slow-marked despite monkeypatching every model construction. Both move to the fast tier; the module-level pytestmark becomes per-test decorators so the 15 real-NNP tests in those two files stay slow. Verified: exactly those two are collected by -m "not slow". Padding invariance gains an ANI2x case at 1e-3, slow-marked. It is not run here: torchani is absent locally and no NNP may be loaded on this box, so CI is what verifies it. Verified: 1302 passed, 9 skipped, 2 xfailed. --- tests/test_SPE.py | 15 ++- tests/test_batchopt.py | 21 +++- tests/test_custom_nnp_contract.py | 88 ++++++++++++++- tests/test_fire_optimizer.py | 76 +++++++++++-- tests/test_model_adapter.py | 88 +++++++++++++++ tests/test_model_factory.py | 2 +- tests/test_optimization_engine_validation.py | 7 +- tests/test_padding_invariance.py | 11 +- tests/test_pipeline_e2e.py | 112 ++++++++++++++++++- tests/test_thermo.py | 20 +++- tests/test_workflow.py | 8 +- 11 files changed, 412 insertions(+), 36 deletions(-) diff --git a/tests/test_SPE.py b/tests/test_SPE.py index d806258c..ddee4b83 100644 --- a/tests/test_SPE.py +++ b/tests/test_SPE.py @@ -7,9 +7,13 @@ # from tests import skip_ani2xt_test skip_ani2xt_test = False -# Mark all tests in this module as slow (single-point energy calculations) -pytestmark = pytest.mark.slow - +# Every real-model test below is marked @pytest.mark.slow individually +# (single-point energy calculations, each loading a real NNP). NOT a +# module-level `pytestmark`: test_calc_spe_uses_model_factory below mocks +# every model-construction call (create_model/EnForce_ANI/pad_from_mols) and +# loads no NNP, so it must run in the fast tier -- a module-level mark would +# have swept it in with everything else regardless of what it actually does. +# # Every calc_spe call below passes use_gpu=False on purpose. calc_spe's # `use_gpu` default is True, and Auto3D 4.0 made "GPU requested but no CUDA # device visible" FATAL rather than a silent CPU fallback @@ -138,6 +142,7 @@ def forward(self, return out['energy'].reshape(-1) +@pytest.mark.slow @pytest.mark.skipif(skip_ani2xt_test, reason="ANI2xt model is not installed.") def test_calc_spe_ani2xt(): #load B97-3c results file @@ -154,6 +159,7 @@ def test_calc_spe_ani2xt(): assert(diff <= 0.01) +@pytest.mark.slow def test_calc_spe_ani2x(): #load wB97X/6-31G* output file path = os.path.join(folder, "tests/files/wb97x_dz.sdf") @@ -170,6 +176,7 @@ def test_calc_spe_ani2x(): print(idx, spe_out, diff) assert(diff <= 0.011) +@pytest.mark.slow def test_calc_spe_aimnet(): path = os.path.join(folder, 'tests/files/cyclooctane.sdf') e_ref = -314.689736079491 @@ -179,6 +186,7 @@ def test_calc_spe_aimnet(): e_out = float(mol.GetProp('E_hartree')) assert(abs(e_out - e_ref) <= 0.01) +@pytest.mark.slow @pytest.mark.skipif(not test_userNNP1, reason="TorchANI is not installed.") def test_calc_spe_userNNP1(): #load wB97X/6-31G* output file @@ -203,6 +211,7 @@ def test_calc_spe_userNNP1(): assert(diff <= 0.011) +@pytest.mark.slow def test_calc_spe_userNNP2(): path = os.path.join(folder, 'tests/files/cyclooctane.sdf') e_ref = -314.689736079491 diff --git a/tests/test_batchopt.py b/tests/test_batchopt.py index 99730bc5..bbabd78c 100644 --- a/tests/test_batchopt.py +++ b/tests/test_batchopt.py @@ -104,8 +104,13 @@ def mock_forward(coords, species, charges, atom_mask=None): energy, forces = model.forward_batched(coords, species, charges) - # Should have called forward multiple times due to batching - assert mock_adapter.forward.call_count >= 1 + # batch_size = max(1, batchsize_atoms // N) = max(1, 10 // 5) = 2 + # molecules per sub-batch; 4 molecules split into chunks of 2 -> the + # adapter must be called exactly twice, not "at least once" (which a + # single unbatched call would also satisfy, defeating the point of a + # forward_batched-specific test). Mirrors test_model_wrapper.py's + # stronger call_count==2 sibling for the same batchsize_atoms/N ratio. + assert mock_adapter.forward.call_count == 2 assert energy.shape == (4,) assert forces.shape == (4, 5, 3) @@ -134,13 +139,19 @@ def test_ensemble_opt_returns_convergence_info(self): result = ensemble_opt(model, coord, numbers, charges, param, torch.device("cpu")) - # Verify new fields are present + # Verify new fields are present, with their actual VALUES: zero force + # on step 1 means fmax (0.0) is at once below opttol (0.01) for both + # structures, so both must be reported converged and neither must + # have been counted as oscillating -- checking only key + # presence/type/length (as before) would pass even if the values + # were transposed, all-False, or a stray increment leaked into + # oscillating_count. assert 'converged_mask' in result, "converged_mask missing from ensemble_opt return" assert 'oscillating_count' in result, "oscillating_count missing from ensemble_opt return" assert isinstance(result['converged_mask'], list) assert isinstance(result['oscillating_count'], list) - assert len(result['converged_mask']) == 2 - assert len(result['oscillating_count']) == 2 + assert result['converged_mask'] == [True, True] + assert result['oscillating_count'] == [0, 0] @pytest.mark.slow diff --git a/tests/test_custom_nnp_contract.py b/tests/test_custom_nnp_contract.py index c04dfe26..e6c722a7 100644 --- a/tests/test_custom_nnp_contract.py +++ b/tests/test_custom_nnp_contract.py @@ -92,6 +92,18 @@ def forward(self, species, coords): return (coords ** 2).sum(dim=(1, 2)) +class TooManyArgsNNP(torch.nn.Module): + """Extra required positional argument -- Auto3D never passes a fourth.""" + + def __init__(self): + super().__init__() + self.coord_pad = 0.0 + self.species_pad = -1 + + def forward(self, species, coords, charges, cutoff): + return (coords ** 2).sum(dim=(1, 2)) * cutoff + + class ExoticNamesNNP(torch.nn.Module): """Right arity, names outside the known vocabulary -- must be accepted.""" @@ -178,7 +190,7 @@ def test_transposed_forward_is_rejected_through_the_adapter(tmp_path): from Auto3D.models.adapter import CustomModelAdapter path = _save(TransposedNNP(), tmp_path, "transposed_adapter.pt") - with pytest.raises(ModelLoadError): + with pytest.raises(ModelLoadError, match="species, coords, charges"): CustomModelAdapter(path, CPU) @@ -209,7 +221,7 @@ def test_transposed_forward_is_rejected_by_input_validation(tmp_path): verbose=False, job_name="", ) - with pytest.raises(ModelLoadError): + with pytest.raises(ModelLoadError, match="species, coords, charges"): check_input(args) @@ -220,6 +232,67 @@ def test_alias_names_in_the_wrong_order_are_rejected(tmp_path): load_custom_nnp(path, CPU) +def _model_with_param_names(species_name, coords_name, charges_name): + """A plain object (no torch.save/pickle needed) whose forward's parameter + NAMES are exactly the given ones, so validate_custom_nnp's order check + (models/contract.py::_classify / _check_forward_signature) sees them. + + Not an nn.Module: this only needs to be introspectable by + inspect.signature, which is all validate_custom_nnp actually uses, and + dynamically-named parameters cannot be produced by a module-level + class (needed elsewhere in this file for torch.save's pickling). + """ + namespace: dict = {} + exec( # noqa: S102 - test-only, fixed trusted template, no user input + f"def forward(self, {species_name}, {coords_name}, {charges_name}):\n" + f" return ({coords_name} ** 2).sum(dim=(1, 2))\n", + namespace, + ) + return type( + "DynamicNamedNNP", + (), + {"coord_pad": 0.0, "species_pad": -1, "forward": namespace["forward"]}, + )() + + +# Full synonym vocabulary from models/contract.py's _SPECIES_NAMES/ +# _COORDS_NAMES/_CHARGES_NAMES, covered at least once each, plus one +# mixed-case ("Numbers"/"Positions"/"Charge") combination to confirm the +# order check case-folds via _classify's ``name.lower()``. +ALIAS_VOCABULARY = [ + ("species", "coords", "charges"), + ("numbers", "positions", "charge"), + ("atomic_numbers", "coordinates", "charge"), + ("atomicnumbers", "coord", "q"), + ("z", "pos", "charges"), + ("elements", "xyz", "charge"), + ("Numbers", "Positions", "Charge"), +] + + +@pytest.mark.parametrize("species_name,coords_name,charges_name", ALIAS_VOCABULARY) +def test_alias_vocabulary_in_correct_order_is_accepted( + species_name, coords_name, charges_name +): + """Every recognized synonym, in the right order, must not be rejected -- + a false rejection here would break a working model that merely spelled + the contract differently.""" + model = _model_with_param_names(species_name, coords_name, charges_name) + validate_custom_nnp(model, "") # must not raise + + +@pytest.mark.parametrize("species_name,coords_name,charges_name", ALIAS_VOCABULARY) +def test_alias_vocabulary_transposed_is_rejected( + species_name, coords_name, charges_name +): + """The same synonyms, transposed (coords first), must still be caught -- + synonyms are not an escape hatch from the order check, across the full + vocabulary, not just the one numbers/positions/charge pair.""" + model = _model_with_param_names(coords_name, species_name, charges_name) + with pytest.raises(ModelLoadError, match="species, coords, charges"): + validate_custom_nnp(model, "") + + def test_missing_both_padding_attributes_are_rejected_at_load(tmp_path): """coord_pad/species_pad are part of the contract; absent, the layers used to disagree on the default, so a silent fallback is worse than a refusal.""" @@ -247,6 +320,17 @@ def test_wrong_arity_is_rejected_at_load(tmp_path): load_custom_nnp(path, CPU) +def test_wrong_arity_too_many_required_args_is_rejected_at_load(tmp_path): + """The ``> 3`` branch: a fourth REQUIRED positional argument is just as + uncallable as the two-argument case above, but exercises the other half + of ``len(positional) < 3 or len(required) > 3`` in + ``models/contract.py::_check_forward_signature``.""" + path = _save(TooManyArgsNNP(), tmp_path, "fourargs.pt") + with pytest.raises(ModelLoadError, match="three positional arguments") as excinfo: + load_custom_nnp(path, CPU) + assert "cutoff" in str(excinfo.value) + + # --- acceptance (a false rejection is a regression) ------------------------- def test_contract_conforming_model_loads_and_runs(tmp_path): diff --git a/tests/test_fire_optimizer.py b/tests/test_fire_optimizer.py index 73a9c4c3..91306fea 100644 --- a/tests/test_fire_optimizer.py +++ b/tests/test_fire_optimizer.py @@ -2,6 +2,7 @@ """Unit tests for the FIRE optimizer module.""" from __future__ import annotations +import pytest import torch from Auto3D.batch_opt.fire_optimizer import FIRE @@ -304,22 +305,73 @@ def test_fire_handles_mixed_convergence(self): assert coord.shape == (4, 5, 3) def test_fire_independent_molecule_tracking(self): - """FIRE should track each molecule's state independently.""" - coord = torch.zeros(3, 5, 3) + """FIRE tracks each molecule's dt/a state independently, driven only + by that molecule's OWN progressing flag -- not by whatever the other + molecules in the batch are doing (fire_optimizer.py's per-molecule + torch.where selects on ``progressing``/``speedup``, but ``a3``, + ``dt``, ``self.a`` etc are shared *tensors* the batch is stepped + through together, so a batch-index mix-up would leak one molecule's + state into another's). + + Protocol (all pure tensor arithmetic; deterministic, no randomness): + Phase 1 -- all 3 molecules push in the same fixed direction, so they + progress together and jointly build up ``Nsteps`` past ``Nmin`` + (bootstrapping requires ALL molecules progressing at least once; see + the ``all_progressing`` branch). Phase 2 -- molecule 0 starts + flipping its own force sign every step. Once misaligned with its own + velocity, a molecule can never re-progress under an alternating + force (each non-progressing step resets v to align with THAT step's + force, and the next step's opposite force is then anti-aligned) -- + so molecule 0 is guaranteed to be "not progressing" for the rest of + phase 2, while molecules 1/2 keep progressing and trigger the + speed-up branch (``past_nmin`` already true from phase 1, + ``all_progressing`` now false because molecule 0 dissents). + """ + n_atoms = 2 + astart = 0.1 + dt_max = 0.1 + coord = torch.zeros(3, n_atoms, 3) optimizer = FIRE(coord) + steady_force = torch.ones(n_atoms, 3) * 0.1 - # Give different molecules different force histories - for i in range(3): - forces = torch.zeros(3, 5, 3) - forces[i] = torch.randn(5, 3) * 0.1 + # Phase 1: all three molecules progress together, bootstrapping + # Nsteps past Nmin (=5) for everyone. + for _ in range(8): + forces = torch.stack([steady_force, steady_force, steady_force]) + coord = optimizer(coord, forces) + assert (optimizer.Nsteps > 5).all(), "phase 1 setup failed to reach Nmin" + + # Phase 2: molecule 0 alternates sign every step (can never progress + # again); molecules 1 and 2 keep pushing steadily. + for step in range(10): + forces = torch.stack( + [ + steady_force if step % 2 == 0 else -steady_force, + steady_force, + steady_force, + ] + ) coord = optimizer(coord, forces) - # Each molecule should have different state - # (at minimum, different velocities) - v_norms = [optimizer.v[i].norm().item() for i in range(3)] - - # They shouldn't all be identical - assert not (v_norms[0] == v_norms[1] == v_norms[2]) + # Molecule 0 never progressed in phase 2, so its mixing parameter + # `a` must be exactly reset to astart every single step -- the + # oscillating molecule's own branch, untouched by its batch-mates. + assert optimizer.a[0].item() == pytest.approx(astart) + + # Molecules 1/2 share an identical force history and so must reach + # an identical (and, since they triggered the speed-up branch, + # strictly smaller-than-astart) mixing parameter -- proving the + # speed-up state is tracked per molecule, not smeared across the + # batch by molecule 0's resets. + assert optimizer.a[1].item() == pytest.approx(optimizer.a[2].item()) + assert optimizer.a[1].item() < astart + + # dt tells the same story from the other side: molecule 0's dt was + # repeatedly shrunk (fdec) by its own non-progress, while 1/2's dt + # saturated at dt_max via their own speed-up. + assert optimizer.dt[1].item() == pytest.approx(dt_max) + assert optimizer.dt[2].item() == pytest.approx(dt_max) + assert optimizer.dt[0].item() < optimizer.dt[1].item() class TestFIRETorchScript: diff --git a/tests/test_model_adapter.py b/tests/test_model_adapter.py index e21fda2a..0777b850 100644 --- a/tests/test_model_adapter.py +++ b/tests/test_model_adapter.py @@ -132,6 +132,47 @@ def test_ani2xt_adapter_creates_model(self): assert adapter.species_pad == -1 assert adapter.coord_pad == 0.0 + def test_ani2xt_adapter_force_sign_with_toy_model(self): + """``forces = -grad`` is duplicated once per adapter in adapter.py: + ``ANI2xtAdapter.forward`` has its own copy, distinct from (and + untested by) ``CustomModelAdapter``'s copy that + ``test_custom_model_adapter_runs`` already checks (audit M32). A sign + bug introduced independently in THIS copy would not be caught by that + test, and ``test_ani2xt_adapter_creates_model`` above never calls + ``.forward`` at all. + + ``BaseModelAdapter.__init__`` is called directly on a bypassed + instance (skipping ``ANI2xtAdapter.__init__``, which imports the real + bundled ANI2xt weights) and handed a toy quadratic model instead -- + same technique ``TestBaseModelAdapter`` already uses for a mock + model, applied here so the REAL ``ANI2xtAdapter.forward`` runs. + Hermetic: no NNP loaded, no torchani import. + """ + from Auto3D.models.adapter import ANI2xtAdapter, BaseModelAdapter + + class _ToyANI2xtModel(torch.nn.Module): + def forward(self, species, coords): + return (coords ** 2).sum(dim=(1, 2)) + + device = torch.device("cpu") + adapter = ANI2xtAdapter.__new__(ANI2xtAdapter) + BaseModelAdapter.__init__( + adapter, _ToyANI2xtModel(), device, coord_pad=0.0, species_pad=-1 + ) + + coords = torch.randn(2, 4, 3) + species = torch.tensor([[0, 1, 2, 3], [0, 1, 2, -1]]) + charges = torch.zeros(2) + energy, forces = adapter.forward(coords, species, charges) + + # _ToyANI2xtModel does not mask padding: E = sum(coords^2) over every + # slot => dE/dx = 2*coords => F = -dE/dx = -2*coords, exactly like + # test_custom_model_adapter_runs's reference calculation. + torch.testing.assert_close( + energy, (coords ** 2).sum(dim=(1, 2)), rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close(forces, -2.0 * coords, rtol=1e-5, atol=1e-6) + class TestANI2xAdapter: """Tests for the ANI2x adapter.""" @@ -150,6 +191,53 @@ def test_ani2x_adapter_creates_model(self): assert adapter.species_pad == -1 assert adapter.coord_pad == 0.0 + def test_ani2x_adapter_force_sign_with_toy_model(self): + """``forces = -grad`` is duplicated again in ``ANI2xAdapter.forward`` + -- a third, separate copy from ``ANI2xtAdapter``'s and + ``CustomModelAdapter``'s (audit M32), also untested by + ``test_custom_model_adapter_runs``. + + The toy model mimics torchani's ``SpeciesEnergies`` return shape (an + object with a ``.energies`` attribute) rather than + ``ANI2xtAdapter``'s plain-tensor return, since ``ANI2xAdapter.forward`` + calls ``self.model((species, coords)).energies`` and multiplies by + ``HARTREE_TO_EV`` -- the toy divides by the same constant first so the + expected force in eV is still the clean ``-2*coords``. Coordinates are + float32 from the start (matching what ``ANI2xAdapter.forward`` casts + to internally) so the adapter's own ``coords.float()`` cast is a + no-op here and cannot be blamed for any looseness in the comparison + (the brainstorm's dtype-cast risk flag for this specific test). + Hermetic: no NNP loaded, no torchani import. + """ + from collections import namedtuple + + from Auto3D.constants import HARTREE_TO_EV + from Auto3D.models.adapter import ANI2xAdapter, BaseModelAdapter + + _SpeciesEnergies = namedtuple("SpeciesEnergies", ["species", "energies"]) + + class _ToyANI2xModel(torch.nn.Module): + def forward(self, species_coords): + species, coords = species_coords + energies = (coords ** 2).sum(dim=(1, 2)) / HARTREE_TO_EV + return _SpeciesEnergies(species, energies) + + device = torch.device("cpu") + adapter = ANI2xAdapter.__new__(ANI2xAdapter) + BaseModelAdapter.__init__( + adapter, _ToyANI2xModel(), device, coord_pad=0.0, species_pad=-1 + ) + + coords = torch.randn(2, 4, 3, dtype=torch.float32) + species = torch.tensor([[1, 6, 7, 8], [1, 6, 7, -1]]) + charges = torch.zeros(2) + energy, forces = adapter.forward(coords, species, charges) + + torch.testing.assert_close( + energy, (coords ** 2).sum(dim=(1, 2)), rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close(forces, -2.0 * coords, rtol=1e-5, atol=1e-6) + class TestCustomModelAdapter: """Tests for the CustomModelAdapter.""" diff --git a/tests/test_model_factory.py b/tests/test_model_factory.py index 80644fc3..c99554a2 100644 --- a/tests/test_model_factory.py +++ b/tests/test_model_factory.py @@ -42,7 +42,7 @@ def __init__(self, *a, **k): raise RuntimeError("unresolvable registry name") monkeypatch.setattr(model_factory, "AIMNet2Adapter", _Boom) - with pytest.raises(Exception): + with pytest.raises(RuntimeError, match="unresolvable registry name"): ModelFactory.create( "totally-not-a-real-model-xyz", device=torch.device("cpu"), diff --git a/tests/test_optimization_engine_validation.py b/tests/test_optimization_engine_validation.py index 8609ad30..760d6b94 100644 --- a/tests/test_optimization_engine_validation.py +++ b/tests/test_optimization_engine_validation.py @@ -153,6 +153,9 @@ def test_n_steps_error_is_not_assertion_error(self): 'nn': None, } - # Should NOT raise AssertionError - with pytest.raises(ValueError): + # Should NOT raise AssertionError, and must be raised for the coord + # shape defect this fixture actually has -- a bare `ValueError` would + # also pass for e.g. an unrelated numbers/charges ValueError, so pin + # the message to the coord/3D guard this fixture is built to hit. + with pytest.raises(ValueError, match="coord.*3D"): n_steps(invalid_state, n=10, opttol=0.01, patience=100) diff --git a/tests/test_padding_invariance.py b/tests/test_padding_invariance.py index 6d99501c..c4640788 100644 --- a/tests/test_padding_invariance.py +++ b/tests/test_padding_invariance.py @@ -31,14 +31,19 @@ class TestPaddingInvariance: # already asserts with atol=1e-2 eV, and ANI2xt's float32 output caps # usable precision at ~float32 ULP (~4e-3 eV) at typical total-energy # magnitudes per src/Auto3D/batch_opt/ANI2xt_no_rep.py:148-155. 1e-6 would - # demand sub-ULP reproducibility and flake on a correct model. + # demand sub-ULP reproducibility and flake on a correct model. ANI2x + # (torchani, periodic-table indexing) shares ANI2xt's float32 output and + # the same -1 species_pad convention, so it gets the same 1e-3 budget + # rather than AIMNet2's looser 1e-2 -- there is no reason to expect it + # tighter than its ANI-family sibling, and no measurement here to justify + # tighter than that either. @pytest.mark.parametrize( "engine, atol", - [("AIMNET", 1e-2), ("ANI2xt", 1e-3)], + [("AIMNET", 1e-2), ("ANI2xt", 1e-3), ("ANI2x", 1e-3)], ) def test_energy_unchanged_when_padded(self, engine, atol, device): """Batching a small molecule alongside a large one must not shift its energy.""" - if engine == "ANI2xt": + if engine in ("ANI2xt", "ANI2x"): pytest.importorskip("torchani") from Auto3D.model_factory import create_model diff --git a/tests/test_pipeline_e2e.py b/tests/test_pipeline_e2e.py index 6f74c97d..a8b7cdf4 100644 --- a/tests/test_pipeline_e2e.py +++ b/tests/test_pipeline_e2e.py @@ -7,7 +7,10 @@ the reconciliation function, exists and is exported and tested with zero production callers (C7). -Slow tier: uses the real aimnet2 registry model on CPU. +Slow tier: uses the real aimnet2 registry model on CPU. NOT a module-level +`pytestmark`, though: `TestClashReliefWarning` below is hermetic (no NNP, no +network) and must run in the fast tier, so every real-pipeline test in this +module is marked `@pytest.mark.slow` individually instead. """ from __future__ import annotations @@ -15,8 +18,7 @@ from rdkit import Chem from Auto3D.config import Auto3DOptions - -pytestmark = pytest.mark.slow +from tests.helpers_pipeline_output import base_molecule_id def _input_ids(smi_path: str) -> set[str]: @@ -32,6 +34,7 @@ def _input_ids(smi_path: str) -> set[str]: class TestInputOutputAccounting: """No input may vanish without being reported.""" + @pytest.mark.slow def test_every_input_is_present_or_reported(self, job_dir): """Each input ID must appear in the output or in a reported failure list. @@ -84,6 +87,7 @@ def test_every_input_is_present_or_reported(self, job_dir): f"{sorted(missing)}" ) + @pytest.mark.slow def test_one_bad_molecule_does_not_remove_the_others(self, job_dir): """A sodium counterion must fail, and must fail alone. @@ -124,8 +128,17 @@ def test_one_bad_molecule_does_not_remove_the_others(self, job_dir): args = Auto3DOptions(path=str(smi), k=1, use_gpu=False, max_confs=2) out = main(args) + # `.split("_")[0]` used to truncate "sodium_acetate" to "sodium" at + # this line, so `"sodium_acetate" not in produced` below was true + # UNCONDITIONALLY -- true whether sodium_acetate correctly failed + # (the intended case) or a regression let it silently succeed and + # reach the output SDF under a name split() would still mangle to + # "sodium". `base_molecule_id` is the pipeline's own id-recovery + # helper (matches what ConformerRanker/decode_ids leave in `_Name` + # by this point), so this now actually depends on whether + # sodium_acetate is or is not in the output. produced = { - m.GetProp("_Name").split("_")[0] + base_molecule_id(m.GetProp("_Name")) for m in Chem.SDMolSupplier(out, removeHs=False) if m is not None } @@ -162,6 +175,7 @@ def test_one_bad_molecule_does_not_remove_the_others(self, job_dir): class TestExitStatus: """Losing molecules must not exit 0.""" + @pytest.mark.slow def test_cli_exits_nonzero_when_molecules_are_missing(self, job_dir): """auto3d run must signal partial failure through its exit code.""" from typer.testing import CliRunner @@ -196,6 +210,7 @@ def test_cli_exits_nonzero_when_molecules_are_missing(self, job_dir): class TestEnergyAndRankingSanity: """Assert on the numbers, not merely that the program ran.""" + @pytest.mark.slow def test_energies_are_negative_and_ordered(self, isolated_input): """E_tot must be negative and ascending within a conformer group.""" from Auto3D.auto3D import main @@ -219,6 +234,7 @@ def test_energies_are_negative_and_ordered(self, isolated_input): f"{base}: conformers are not energy-ordered: {energies}" ) + @pytest.mark.slow def test_top_k_returns_distinct_conformers(self, isolated_input): """k=3 must yield at most 3 per molecule, and the first is the minimum.""" from Auto3D.auto3D import main @@ -238,3 +254,91 @@ def test_top_k_returns_distinct_conformers(self, isolated_input): for base, energies in groups.items(): assert len(energies) <= 3, f"{base}: k=3 but got {len(energies)}" assert energies[0] == min(energies), f"{base}: first is not the minimum" + + +class TestClashReliefWarning: + """RDKitIsomer's serial embedding path (isomer_engine.py's + ``_run_serial_embedding``) must warn, once, when EVERY embedded conformer + of a species is rejected by clash relief -- the species then silently + vanishes from the output with no other trace. Before this test, + ``grep -rn "produced no conformers after clash relief" tests/`` returned + nothing: this is distinct from (and untested by) the parallel-embedding + path's own version of the same warning in ``isomers/parallel_embed.py``, + which ``test_workflow.py`` already covers. + + Hermetic: no NNP, no network. ``relieve_clash`` itself is monkeypatched, + so this exercises only isomer_engine.py's warn-and-continue behavior + around it, not the clash-relief force field logic. + """ + + def test_species_with_no_surviving_conformer_is_warned_and_a_sibling_is_not( + self, tmp_path, caplog, monkeypatch + ): + import logging + + import Auto3D.isomer_engine as isomer_engine_mod + from Auto3D.isomer_engine import RDKitIsomer + + smi = tmp_path / "in.smi" + # "bad_mol" (methane) will have every conformer rejected below; + # "good_mol" (ethanol) is the companion the guard must NOT warn about + # (the over-fire check the spec's verification standard asks for). + smi.write_text("C\tbad_mol\nCCO\tgood_mol\n") + + job_dir = tmp_path / "job" + job_dir.mkdir() + + engine = RDKitIsomer( + smi=str(smi), + smiles_enumerated=str(tmp_path / "enumerated.smi"), + smiles_enumerated_reduced=str(tmp_path / "enumerated_reduced.smi"), + smiles_hashed=str(tmp_path / "hashed.smi"), + enumerated_sdf=str(tmp_path / "enumerated.sdf"), + job_name=str(job_dir), + max_confs=1, + threshold=0.3, + np=1, + flipper=False, + ) + + # relieve_clash(mol, conf_id) never sees the species name -- only the + # embedded RDKit Mol -- so identify "the methane-derived molecule" + # the only way the stub can: by its (distinctive) atom count after + # AddHs. Bypassing the real force-field logic entirely keeps this + # test hermetic and removes any dependence on ETKDG's random seed + # actually producing a clashing geometry. + bad_atom_count = Chem.AddHs(Chem.MolFromSmiles("C")).GetNumAtoms() + + def fake_relieve_clash(mol, conf_id): + return mol.GetNumAtoms() != bad_atom_count + + monkeypatch.setattr( + isomer_engine_mod, "relieve_clash", fake_relieve_clash + ) + + with caplog.at_level(logging.WARNING): + out = engine.run() + + warnings = [ + r.message for r in caplog.records if r.levelno == logging.WARNING + ] + clash_warnings = [ + m for m in warnings if "produced no conformers after clash relief" in m + ] + assert len(clash_warnings) == 1, ( + f"expected exactly one clash-relief warning, got {clash_warnings}" + ) + assert "bad_mol" in clash_warnings[0] + # The over-fire check: a sibling molecule in the same batch that DID + # survive clash relief must not be named in any such warning. + assert "good_mol" not in clash_warnings[0] + assert not any("good_mol" in m for m in clash_warnings) + + # bad_mol is absent from the output entirely; good_mol made it through. + produced = { + m.GetProp("_Name") + for m in Chem.SDMolSupplier(out, removeHs=False) + if m is not None + } + assert not any(name.startswith("bad_mol") for name in produced), produced + assert any(name.startswith("good_mol") for name in produced), produced diff --git a/tests/test_thermo.py b/tests/test_thermo.py index e2b7429d..71ea6b0e 100644 --- a/tests/test_thermo.py +++ b/tests/test_thermo.py @@ -14,9 +14,13 @@ write_perturbed_sdf, ) -# Mark all tests in this module as slow (thermodynamic calculations) -pytestmark = pytest.mark.slow - +# Every real-model test below is marked @pytest.mark.slow individually +# (thermodynamic calculations, each loading a real NNP). NOT a module-level +# `pytestmark`: test_model_name2model_calculator_uses_factory below patches +# both create_model and EnForce_ANI and loads no NNP, so it must run in the +# fast tier -- a module-level mark would have swept it in regardless (its +# test_SPE.py twin, test_calc_spe_uses_model_factory, had the same defect). +# # Every opt_geometry/calc_thermo call below passes use_gpu=False on purpose. # Both default to use_gpu=True, and Auto3D 4.0 made "GPU requested but no CUDA # device visible" FATAL rather than a silent CPU fallback @@ -252,6 +256,7 @@ def assert_thermo_record(mol, *, reference_G=None, reference_H=None): ) +@pytest.mark.slow def test_calc_thermo_aimnet(): """AIMNET thermochemistry for cyclooctane against a wB97m-D4/Def2-TZVPP run. @@ -273,6 +278,7 @@ def test_calc_thermo_aimnet(): except OSError: pass +@pytest.mark.slow def test_vib_hessian_includes_external_dispersion(): """Regression guard: the AIMNET vibrational Hessian must run the full energy pipeline (external D3 dispersion + Coulomb), not the bare aimnet nn.Module. @@ -371,6 +377,7 @@ def _perturbed_DA(tmp_path) -> tuple[str, list]: return write_perturbed_sdf(source, tmp_path / "DA.sdf", DA_EXPANSION) +@pytest.mark.slow def test_opt_geometry1(tmp_path): """ANI2x relaxes a displaced geometry and annotates it correctly.""" path, inputs = _perturbed_DA(tmp_path) @@ -378,6 +385,7 @@ def test_opt_geometry1(tmp_path): assert_opt_geometry_output(out, input_mols=inputs, moved_at_least=DA_MIN_RELAXATION, label="ANI2x") +@pytest.mark.slow def test_opt_geometry2(tmp_path): """ANI2xt relaxes a displaced geometry and annotates it correctly.""" path, inputs = _perturbed_DA(tmp_path) @@ -385,6 +393,7 @@ def test_opt_geometry2(tmp_path): assert_opt_geometry_output(out, input_mols=inputs, moved_at_least=DA_MIN_RELAXATION, label="ANI2xt") +@pytest.mark.slow def test_opt_geometry3(tmp_path): """AIMNet2 relaxes a displaced geometry and annotates it correctly.""" path, inputs = _perturbed_DA(tmp_path) @@ -393,6 +402,7 @@ def test_opt_geometry3(tmp_path): moved_at_least=DA_MIN_RELAXATION, label="AIMNET") +@pytest.mark.slow def test_opt_geometry_with_patience_and_batchsize(): """Test opt_geometry with explicit patience and batchsize_atoms parameters.""" path = os.path.join(folder, "tests/files/DA.sdf") @@ -411,6 +421,7 @@ def test_opt_geometry_with_patience_and_batchsize(): except OSError: pass +@pytest.mark.slow @pytest.mark.skipif(not test_userNNP1, reason="TorchANI is not installed.") def test_opt_geometry4(tmp_path): """A scripted custom NNP relaxes a displaced geometry through opt_geometry.""" @@ -426,6 +437,7 @@ def test_opt_geometry4(tmp_path): moved_at_least=DA_MIN_RELAXATION, label="scripted userNNP1") +@pytest.mark.slow def test_opt_geometry5(tmp_path): """An eager AIMNet2-backed custom NNP relaxes a displaced geometry.""" path, inputs = _perturbed_DA(tmp_path) @@ -441,6 +453,7 @@ def test_opt_geometry5(tmp_path): label="eager userNNP2") +@pytest.mark.slow @pytest.mark.skipif(not test_userNNP1, reason="TorchANI is not installed.") def test_calc_thermo_userNNP1(): #load wB97m-D4/Def2-TZVPP output file @@ -478,6 +491,7 @@ def test_calc_thermo_userNNP1(): pass +@pytest.mark.slow def test_calc_thermo_userNNP2(): #load wB97m-D4/Def2-TZVPP output file path = os.path.join(folder, "tests/files/cyclooctane.sdf") diff --git a/tests/test_workflow.py b/tests/test_workflow.py index 9f4d533f..5fd996c2 100644 --- a/tests/test_workflow.py +++ b/tests/test_workflow.py @@ -297,7 +297,13 @@ def test_optimizer_handles_empty_file(self, tmp_path, caplog, monkeypatch): with caplog.at_level(logging.WARNING): optimizer.run() - assert "empty" in caplog.text + # Pin the exact guard that fired: the empty-file message must be + # distinguishable from the missing-file message above ("does not + # exist"), which is also a file literally named "empty.sdf" would + # trivially satisfy a bare "empty" in caplog.text check without ever + # proving the *empty-file* branch (not the missing-file branch) ran. + assert f"Input file {empty_sdf} is empty." in caplog.text + assert "does not exist" not in caplog.text def test_workers_importable_from_workflow_workers(): From d3ce08b40076cc8f2c35f6efd8a4571468f277b9 Mon Sep 17 00:00:00 2001 From: isayev Date: Tue, 4 Aug 2026 02:06:50 -0400 Subject: [PATCH 5/6] refactor!: one model contract, and one authoritative config schema Two independent clusters that had to land together because both touch auto3D.py's call sites. ## One model contract ModelAdapter moves from models/adapter.py into models/contract.py beside CustomNNP, 40 lines apart, so the deliberate argument-order inversion between them is structural rather than a comment: CustomNNP is forward(species, coords, charges) -> energies and is the public custom-NNP contract, while ModelAdapter is forward(coords, species, charges) -> (energies, forces) and is internal. Confusing the two is the failure this layout is meant to prevent. REQUIRED_ATTRIBUTES is now derived from the Protocol's own annotations rather than hand-maintained beside it, so the two cannot drift. That makes those annotations load-bearing for TorchScript archives, since contract.py skips the signature check for them and is the only gate; the code says so. CustomNNP loses @runtime_checkable. It was never used, and it would not have helped: isinstance against a Protocol only checks attribute presence, so it cannot see Module.forward's stub, which the validator already handles by hand. isinstance now raises TypeError pointing at validate_custom_nnp. batch_opt no longer imports model_factory. The adapter is injected instead, constructed inside the worker so it never crosses a spawn boundary, and batch_opt/species.py is deleted in favour of models/species.py, which puts the species convention next to the model that defines it. Seven copy-pasted MLP definitions in ANI2xt_no_rep.py collapse to a WIDTHS table and a factory, 75 lines to 9. Verified rather than assumed: a test loads the shipped ani2xt_no_repulsion.pt into both the old seven-block ModuleList and the new one and compares. All 56 tensors satisfy torch.equal with identical key sets, because ModuleList state_dict keys are positional and never derived from Python variable names. The old source declared S_network before F_network but placed F before S in the list, which was harmless only because F, S and Cl share widths. Three adapters get their own dtype-preserving energy(). Defining it as forward(...)[0] would have silently downgraded an fp64 Hessian to fp32, because ANI2xAdapter and CustomModelAdapter both call coords.float() and ANI2xt additionally calls requires_grad_(True), which raises on the non-leaf tensor an autograd Hessian supplies. No error, just quietly worse thermochemistry. Four tests pin it. ## One authoritative config schema check_valid_configuration took ten keyword arguments and carried a third set of defaults, including a literal opt_steps=2000. It now takes an Auto3DOptions, which is the single source of truth, and the two byte-identical ten-kwarg marshalling blocks in auto3D.py and workflow.py become one line each. Engine choices move to one ENGINE_CHOICES table. opt_steps had two different minimums: FIELD_BOUNDS declared 1 while validation.py hand-wrote >= 10 twice. 10 is correct and wins. 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; below 10 there is no early exit, no progress and no reporting, and FIRE needs several steps to build velocity. So opt_steps < 10 returned an unconverged structure labelled optimized. Loosening the bound to 1 would have accepted that; tightening only moves an existing refusal earlier. The legacy `auto3d cfg.yaml` path had three shape guards the `run -c` path did not, so an empty file, a non-mapping top level or a syntax error exited 1 "Unexpected Error" instead of 2 ConfigurationError. It now goes through load_yaml_config, and the banner is printed after validation so an unrunnable config is never announced as running. The test that should have caught this claimed in its docstring to use "the exact construction _run_legacy_yaml uses" and then re-implemented it inline. That replica is why the missing guards were invisible; it now calls the real function. A test walks CLIConfig.model_fields metadata and rejects any pydantic constraint on a field whose bound lives in FIELD_BOUNDS, so the duplicated constraint set that had to be backed out once cannot come back. Verified: 1349 passed, 9 skipped, 1 xfailed, ruff clean. The remaining xfail is the E_tot filter divergence, owned by a later cluster; the opt_steps pin is resolved and its marker removed. Every fix mutation-verified by reverting it and confirming the named test goes red. --- parameters.yaml | 8 + src/Auto3D/ASE/geometry.py | 10 +- src/Auto3D/ASE/thermo.py | 2 +- src/Auto3D/SPE.py | 7 +- src/Auto3D/auto3D.py | 22 +- src/Auto3D/auto3Dcli.py | 83 +++--- src/Auto3D/batch_opt/ANI2xt_no_rep.py | 133 ++++----- src/Auto3D/batch_opt/batchopt.py | 53 ++-- src/Auto3D/batch_opt/model_wrapper.py | 43 ++- src/Auto3D/batch_opt/padding.py | 36 ++- src/Auto3D/batch_opt/species.py | 65 ----- src/Auto3D/cli/commands/models.py | 11 +- src/Auto3D/cli/config_schema.py | 136 +++++---- src/Auto3D/config.py | 76 ++++- src/Auto3D/model_factory.py | 18 +- src/Auto3D/models/__init__.py | 16 +- src/Auto3D/models/adapter.py | 203 +++++++++++--- src/Auto3D/models/contract.py | 253 +++++++++++++++-- src/Auto3D/models/species.py | 100 +++++++ src/Auto3D/utils/validation.py | 111 ++++---- src/Auto3D/workflow.py | 13 +- src/Auto3D/workflow_workers.py | 15 +- tests/helpers_adapter.py | 113 ++++++++ tests/test_SPE.py | 6 +- tests/test_adapter_atom_mask.py | 25 +- tests/test_batchopt.py | 173 +++++++----- tests/test_cli_config_schema.py | 385 ++++++++++++++++++++++++-- tests/test_cli_exit_codes.py | 7 +- tests/test_cli_property_commands.py | 28 +- tests/test_config.py | 76 +++-- tests/test_config_parity.py | 4 +- tests/test_custom_nnp_contract.py | 167 +++++++++-- tests/test_durability.py | 8 +- tests/test_e_tot_units.py | 45 +-- tests/test_isomer_engine_hardening.py | 2 +- tests/test_legacy_yaml_parity.py | 144 ++++++++++ tests/test_model_adapter.py | 293 +++++++++++++++++++- tests/test_model_factory.py | 34 ++- tests/test_model_preflight.py | 15 +- tests/test_model_wrapper.py | 78 +++++- tests/test_padding.py | 92 +++++- tests/test_padding_invariance.py | 19 +- tests/test_species_conversion.py | 17 +- tests/test_species_module.py | 95 ++++++- tests/test_thermo_helpers.py | 7 +- tests/test_thermo_transition_state.py | 6 +- tests/test_utils_validation.py | 144 +++++----- tests/test_validation.py | 38 ++- tests/test_workflow.py | 41 +-- 49 files changed, 2701 insertions(+), 775 deletions(-) delete mode 100644 src/Auto3D/batch_opt/species.py create mode 100644 src/Auto3D/models/species.py create mode 100644 tests/helpers_adapter.py create mode 100644 tests/test_legacy_yaml_parity.py 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..c5f61fad 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): 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): From 5555e650d46c208d9f71f61e47f46017f62b4fef Mon Sep 17 00:00:00 2001 From: isayev Date: Tue, 4 Aug 2026 02:10:27 -0400 Subject: [PATCH 6/6] docs: correct the signatures the model-contract change left stale Three places described code that no longer exists: config.py's docstring example still called optimizing positionally with a model, and the migration guide still showed pad_from_mols taking a model name and to_model_species taking an engine string. The migration guide's warning about optimizing binding a stray positional argument to progress_cb is replaced rather than edited. That hazard is gone: everything after out_f is keyword-only now, so the call raises TypeError instead of silently accepting a wrong-typed callback that n_steps would then swallow in `except Exception: pass`. The section records what the hazard was, since it explains why the parameters are keyword-only. Adds the wave-2 breaking changes to the CHANGELOG: the adapter injection and its signature table, the opt_steps minimum with the reasoning for choosing 10, check_valid_configuration's new parameter, and the legacy YAML exit-code correction. --- CHANGELOG.md | 52 +++++++++++++++++++++++++++++++++++ docs/source/migration-4.0.rst | 39 ++++++++++++++++++-------- src/Auto3D/config.py | 2 +- 3 files changed, 80 insertions(+), 13 deletions(-) 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/src/Auto3D/config.py b/src/Auto3D/config.py index c5f61fad..60683018 100644 --- a/src/Auto3D/config.py +++ b/src/Auto3D/config.py @@ -465,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