From 80e244b87f86841b5d44bd0a2b8712fda58a24be Mon Sep 17 00:00:00 2001 From: isayev Date: Mon, 3 Aug 2026 18:01:00 -0400 Subject: [PATCH] refactor!: retire EnForce_ANI's type-switched second parameter BREAKING CHANGE: EnForce_ANI's signature is now (model_adapter, batchsize_atoms=16384). The second parameter used to be `name_or_batchsize: str | int | None`, switching between a model name and a batch size; passing a string warned it would be "removed in Auto3D v2.0". The package reached 3.0.0 with it still there -- two majors past its own removal notice -- and no caller in src/ ever passed one. `_legacy_forward`, 55 lines dispatching on `self.name`, is deleted with it. Verified before acting, per the practice this backlog now requires: the three production callers are batchopt.py:319 (positional int), ASE/thermo.py:590 and SPE.py:148 (adapter only). None uses the string form, and the positional int call keeps working unchanged against the two-parameter signature. Removing the union on its own would have left a trap for exactly the callers the removal affects: `EnForce_ANI(adapter, "AIMNET")` would assign a string to batchsize_atoms and surface much later inside batching as a comparison error naming neither the parameter nor the removal. It now raises TypeError on the spot, saying what the parameter is for and where to get an adapter. Two tests asserted `_use_legacy_forward is False` while actually testing batch-size handling; they keep their real assertions and lose the one about a deleted attribute. The test of the removed feature is replaced by one asserting the old call shape is now rejected. Mutation-verified: dropping the guard fails it. Verified: 1273 passed, 9 skipped; ruff clean. --- CHANGELOG.md | 14 +++ .../follow-ups-after-4.0.0-remediation.md | 2 +- src/Auto3D/batch_opt/model_wrapper.py | 109 +++--------------- tests/test_model_wrapper.py | 36 +++--- 4 files changed, 54 insertions(+), 107 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7e2b1878..5eb8c64c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1279,6 +1279,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 into every CLI command (`run`, `energy`, `optimize`, `thermo`, `tautomers`, `models test`, and the legacy YAML path). +- **`EnForce_ANI`'s type-switched second parameter is gone.** It was + `name_or_batchsize: str | int | None`, switching between a model name (the + pre-adapter API) and a batch size, and passing a string warned that it would be + *"removed in Auto3D v2.0"*. The package reached 3.0.0 with it still in place, + two majors past its own removal notice, and no caller in `src/` ever passed one. + The signature is now `EnForce_ANI(model_adapter, batchsize_atoms=16384)` and + `_legacy_forward` (55 lines dispatching on `self.name`) is deleted. + + **Migration:** build an adapter with `Auto3D.model_factory.create_model` and pass + it as the first argument. A string in the second position now raises `TypeError` + naming the parameter — without that guard, removing the union would have silently + assigned a model name to `batchsize_atoms` and failed much later inside batching, + as a comparison error mentioning neither the parameter nor the removal. + - **`use_parallel_embedding` is now reachable.** Parallel conformer embedding existed as a constructor argument on the isomer engine with no route from `Auto3DOptions`, so no `main()` or `smiles2mols` run could turn it on and the diff --git a/docs/superpowers/follow-ups-after-4.0.0-remediation.md b/docs/superpowers/follow-ups-after-4.0.0-remediation.md index df633729..a1645407 100644 --- a/docs/superpowers/follow-ups-after-4.0.0-remediation.md +++ b/docs/superpowers/follow-ups-after-4.0.0-remediation.md @@ -121,7 +121,7 @@ present and dead: | **already deleted** by earlier phases | `pad_molecular_batch`, `cli/progress` `create_progress`, `IsomerProgressCallback` | | **not dead — M53 wrong** | `utils/stereo_check` (6 live uses), `cli/results` `FailedMolecule` + `print_failures` (live from `run.py` since the C6/C7 reconciliation work — the finding's "run.py admits failures is always []" no longer holds), `ASE/thermo` `mol2atoms`, `STANDARD_PRESSURE`, `isomers/parallel_embed` | | **genuinely dead — deleted** | `utils_file.py` (whole module), `count_from_output`, `encode_smiles`, `decode_smiles`, `housekeeping_helper`, and 3 constants (`BOND_STRETCH_TOLERANCE`, `COLLISION_THRESHOLD`, `SUPPORTED_MODELS`) | -| **unresolved — line numbers stale** | `exceptions.py` "4 classes never raised" (line 41 is `OptimizationError`, raised 3x, so the cited lines no longer point at what the finding describes); `model_wrapper`'s legacy `name` API; `ASE/thermo` `model_name` param | +| **unresolved — line numbers stale** | `exceptions.py` "4 classes never raised" (line 41 is `OptimizationError`, raised 3x, so the cited lines no longer point at what the finding describes); `ASE/thermo` `model_name` param (the `model_wrapper` legacy `name` API is **DONE 2026-08-03**) | Net: **256 lines removed from `src/`, 167 from `tests/`** — not the ~450 of `src/` the finding claimed, because a third of it was gone and half of the rest is alive. diff --git a/src/Auto3D/batch_opt/model_wrapper.py b/src/Auto3D/batch_opt/model_wrapper.py index 2523a3f8..b90f2433 100644 --- a/src/Auto3D/batch_opt/model_wrapper.py +++ b/src/Auto3D/batch_opt/model_wrapper.py @@ -6,14 +6,12 @@ """ from __future__ import annotations -import warnings from typing import TYPE_CHECKING import torch import torch.nn as nn from Auto3D.exceptions import OptimizationError -from Auto3D.utils import hartree2ev if TYPE_CHECKING: from Auto3D.model_factory import BaseModelAdapter @@ -26,9 +24,8 @@ class EnForce_ANI(nn.Module): for calculating energies and forces. Args: - model_adapter: A model adapter implementing the forward(coords, species, charges) interface, - or a raw model (for backward compatibility). - name_or_batchsize: Either a string name (deprecated old API) or an int batchsize_atoms. + model_adapter: A model adapter implementing the + forward(coords, species, charges) interface. batchsize_atoms: Maximum number of atoms that can be handled in one batch. Examples: @@ -45,44 +42,33 @@ class EnForce_ANI(nn.Module): def __init__( self, model_adapter: BaseModelAdapter, - name_or_batchsize: str | int | None = None, batchsize_atoms: int = 1024 * 16, ) -> None: """Initialize EnForce_ANI wrapper. Args: model_adapter: A model adapter implementing the forward interface. - name_or_batchsize: For backward compatibility - either model name (deprecated) - or batchsize_atoms as int. batchsize_atoms: Maximum number of atoms per batch (default: 16384). + + 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 + package reached 3.0.0 with it still in place, and no caller in ``src/`` + ever passed one. Removed, so the parameter has one meaning. """ super().__init__() - # Handle backward compatibility - if isinstance(name_or_batchsize, str): - # Old API: EnForce_ANI(model, name, batchsize_atoms) - warnings.warn( - "Passing 'name' to EnForce_ANI is deprecated and will be removed in Auto3D v2.0. " - "Use model adapters from Auto3D.model_factory instead.", - DeprecationWarning, - stacklevel=2, + # 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 + # the spot, naming what the parameter is now for. + if not isinstance(batchsize_atoms, int) or isinstance(batchsize_atoms, bool): + raise TypeError( + "EnForce_ANI's second parameter is batchsize_atoms (an int), got " + f"{batchsize_atoms!r}. The model-name form was removed in 3.0.0; " + "build an adapter with Auto3D.model_factory.create_model instead." ) - self.add_module("ani", model_adapter) - self.model = model_adapter - self.name = name_or_batchsize - self.batchsize_atoms = batchsize_atoms - self._use_legacy_forward = True - elif isinstance(name_or_batchsize, int): - # New API with explicit batchsize: EnForce_ANI(model_adapter, batchsize_atoms) - self.model = model_adapter - self.batchsize_atoms = name_or_batchsize - self.name = None - self._use_legacy_forward = False - else: - # New API: EnForce_ANI(model_adapter) or EnForce_ANI(model_adapter, None, batchsize) - self.model = model_adapter - self.batchsize_atoms = batchsize_atoms - self.name = None - self._use_legacy_forward = False + self.model = model_adapter + self.batchsize_atoms = batchsize_atoms def forward( self, @@ -120,65 +106,8 @@ def forward( Tuple of (energies, forces) where energies has shape (B,) and forces has shape (B, N, 3). """ - if self._use_legacy_forward: - return self._legacy_forward(coord, numbers, charges) return self.model.forward(coord, numbers, charges, atom_mask=atom_mask) - def _legacy_forward( - self, - coord: torch.Tensor, - numbers: torch.Tensor, - charges: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Legacy forward implementation for backward compatibility. - - .. deprecated:: 1.0 - This method is deprecated and will be removed in Auto3D v2.0. - Use model adapters from :mod:`Auto3D.model_factory` instead. - - Handles raw models that were passed with the old API. - - Note: Cannot use torch.inference_mode() because force calculation - requires computing gradients via torch.autograd.grad(). Model parameters - are already frozen (requires_grad=False), but coordinates must have - requires_grad=True for force computation. - - Args: - coord: Coordinates for all input structures. - numbers: The periodic numbers for all atoms. - charges: Molecular charges. - - Returns: - Tuple of (energies, forces). - """ - warnings.warn( - "_legacy_forward is deprecated and will be removed in Auto3D v2.0. " - "Use model adapters from Auto3D.model_factory instead.", - DeprecationWarning, - stacklevel=2, - ) - if self.name == "AIMNET": - d = self.ani(dict(coord=coord, numbers=numbers, charge=charges)) - e = d["energy"].to(torch.double) - f = d["forces"] - elif self.name == "ANI2xt": - e = self.ani(numbers, coord) - # create_graph=False avoids building second-order gradient graph - g = torch.autograd.grad([e.sum()], [coord], create_graph=False)[0] - f = -g - elif self.name == "ANI2x": - e = self.ani((numbers, coord)).energies - e = e * hartree2ev # ANI2x output energy unit is Hartree; convert to eV - # create_graph=False avoids building second-order gradient graph - g = torch.autograd.grad([e.sum()], [coord], create_graph=False)[0] - f = -g - else: - # user NNP that was loaded from a file - e = self.ani(numbers, coord, charges) - # create_graph=False avoids building second-order gradient graph - g = torch.autograd.grad([e.sum()], [coord], create_graph=False)[0] - f = -g - return e, f def forward_batched( self, diff --git a/tests/test_model_wrapper.py b/tests/test_model_wrapper.py index 69ad8a45..5eb03f1f 100644 --- a/tests/test_model_wrapper.py +++ b/tests/test_model_wrapper.py @@ -47,7 +47,6 @@ def test_enforce_ani_forward_with_batchsize_kwarg(self): wrapper = EnForce_ANI(mock_adapter, batchsize_atoms=512) assert wrapper.batchsize_atoms == 512 - assert wrapper._use_legacy_forward is False def test_enforce_ani_forward_with_int_second_arg(self): """EnForce_ANI should accept int as second argument for batchsize.""" @@ -60,7 +59,6 @@ def test_enforce_ani_forward_with_int_second_arg(self): wrapper = EnForce_ANI(mock_adapter, 256) assert wrapper.batchsize_atoms == 256 - assert wrapper._use_legacy_forward is False class TestEnForceANIForwardBatched: @@ -164,20 +162,6 @@ def mock_forward(coords, species, charges, atom_mask=None): class TestEnForceANIBackwardCompatibility: """Tests for backward compatibility with legacy API.""" - def test_enforce_ani_legacy_api_emits_deprecation_warning(self): - """Using string name should emit deprecation warning.""" - # Need a real nn.Module for the legacy API since it uses add_module() - mock_model = torch.nn.Linear(1, 1) - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - wrapper = EnForce_ANI(mock_model, "AIMNET", 1024) - - assert len(w) == 1 - assert issubclass(w[0].category, DeprecationWarning) - assert "deprecated" in str(w[0].message).lower() - assert wrapper._use_legacy_forward is True - assert wrapper.name == "AIMNET" def test_enforce_ani_import_from_batchopt(self): """EnForce_ANI should be importable from batchopt for backward compatibility.""" @@ -243,3 +227,23 @@ def forward(self, coord, numbers, charges, atom_mask=None): e, f = wrapper.forward_batched(coord, numbers, charges) assert e.shape == (2,) and torch.isfinite(e).all() assert f.shape == coord.shape + + +def test_a_model_name_in_the_batchsize_slot_is_rejected(): + """The removed API's shape must fail loudly, not become a bad batch size. + + Until 3.0.0 the second parameter was `name_or_batchsize: str | int | None`, + type-switched between a model name and a batch size, and passing a string + warned it would be "removed in Auto3D v2.0". The package reached 3.0.0 with it + still there and no caller in `src/` ever passing one, so it is gone. + + With the union removed and nothing else added, `EnForce_ANI(adapter, "AIMNET")` + would have assigned a string to `batchsize_atoms` and failed much later inside + batching, as a comparison error naming neither the parameter nor the removal. + """ + import pytest + + from Auto3D.batch_opt.model_wrapper import EnForce_ANI + + with pytest.raises(TypeError, match="batchsize_atoms"): + EnForce_ANI(MagicMock(), "AIMNET")