Merge pull request #18 from erwallace/v0.0.10

v0.0.10

Author: erwallace

README.md

@@ -134,12 +134,6 @@ class MyCalculator(Calculator):
         forces = torch.zeros_like(batch.pos, device=self.device)
         # ... fill forces from your model ...
         return energies, forces
-
-    def to_atomic_data():
-        pass
-
-    def from_atomic_data():
-        pass
 ```
 
 ### Data Containers

src/neural_optimiser/calculators/__init__.py

@@ -1,5 +1,6 @@
 from ._fairchem import FAIRChemCalculator
 from ._mace import MACECalculator
+from ._mmff94 import MMFF94Calculator
 from ._rand import RandomCalculator
 
-__all__ = ["MACECalculator", "FAIRChemCalculator", "RandomCalculator"]
+__all__ = ["MACECalculator", "FAIRChemCalculator", "RandomCalculator", "MMFF94Calculator"]

src/neural_optimiser/calculators/_mmff94.py

@@ -0,0 +1,109 @@
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from ase.units import eV, kcal, mol
+from rdkit import Chem
+from rdkit.Chem import AllChem, rdDetermineBonds
+from torch_geometric.data import Batch, Data
+
+from neural_optimiser.calculators.base import Calculator
+
+KCAL_PER_MOL_TO_EV = kcal / mol / eV
+
+
+class MMFF94Calculator(Calculator):
+    """Calculator using RDKit's implementation of the Merck Molecular Force Field (MMFF94(s))."""
+
+    def __init__(
+        self,
+        MMFFGetMoleculeProperties: Mapping | None = None,
+        MMFFGetMoleculeForceField: Mapping | None = None,
+    ):
+        self.mol_prop = dict(MMFFGetMoleculeProperties or {})
+        self.mol_ff = dict(MMFFGetMoleculeForceField or {})
+        self.bond_info = None
+        self._cache_key = None
+
+    def __repr__(self) -> str:
+        parts = []
+        if self.mol_prop:
+            parts.extend(f"{k}={v}" for k, v in self.mol_prop.items())
+        if self.mol_ff:
+            parts.extend(f"{k}={v}" for k, v in self.mol_ff.items())
+        args = ", ".join(parts)
+        return f"MMFFCalculator({args})" if args else "MMFFCalculator()"
+
+    def _ensure_single_conformer(self, batch: Data | Batch) -> None:
+        """Ensure that the batch contains only a single conformer."""
+        if isinstance(batch, Batch) and batch.batch.unique().size(0) != 1:
+            raise ValueError("MMFFCalculator only supports single-conformer batches.")
+
+    def _to_mol(self, batch: Data | Batch) -> Chem.RWMol:
+        """Convert a single-conformer batch to an RDKit RWMol."""
+        if isinstance(batch, Batch):
+            return Chem.RWMol(batch.to_rdkit()[0])
+        return Chem.RWMol(batch.to_rdkit())
+
+    def _get_charge(self, mol: Chem.Mol) -> int:
+        """Get the formal charge of the molecule from its properties, defaulting to 0."""
+        return int(mol.GetProp("charge")) if mol.HasProp("charge") else 0
+
+    def _molecule_key(self, mol: Chem.Mol) -> str:
+        """Generate a unique key for the molecule based on its SMILES or atom types."""
+        if mol.HasProp("smiles"):
+            s = mol.GetProp("smiles")
+            if s:
+                return s
+        # Fallback to SMILES; as a last resort, atom sequence signature
+        try:
+            return Chem.MolToSmiles(mol, canonical=True)
+        except Exception:
+            return ",".join(str(a.GetAtomicNum()) for a in mol.GetAtoms())
+
+    def _prepare_mol(self, mol: Chem.Mol) -> Chem.Mol:
+        """Determine or restore bonds, then sanitize. Uses cached bond topology per molecule key."""
+        charge = self._get_charge(mol)
+        new_key = self._molecule_key(mol)
+
+        if new_key != self._cache_key:
+            rdDetermineBonds.DetermineBonds(mol, charge=charge)
+            self.bond_info = [
+                (b.GetBeginAtomIdx(), b.GetEndAtomIdx(), b.GetBondType()) for b in mol.GetBonds()
+            ]
+            self._cache_key = new_key
+        else:
+            if self.bond_info:
+                for begin, end, bond_type in self.bond_info:
+                    mol.AddBond(begin, end, bond_type)
+
+        Chem.SanitizeMol(mol)
+        return mol
+
+    def _build_forcefield(self, mol: Chem.Mol) -> Any:
+        """Build the MMFF force field for the given molecule."""
+        mp = AllChem.MMFFGetMoleculeProperties(mol, **self.mol_prop)
+        ff = AllChem.MMFFGetMoleculeForceField(mol, mp, **self.mol_ff)
+        return ff
+
+    def _calculate(self, batch: Data | Batch) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute energies and forces for a batch of conformers using MMFF94."""
+        self._ensure_single_conformer(batch)
+        mol = self._to_mol(batch)
+        mol = self._prepare_mol(mol)
+
+        ff = self._build_forcefield(mol)
+        energy = torch.tensor([ff.CalcEnergy()])
+        grad = (
+            torch.Tensor(ff.CalcGrad()) * -KCAL_PER_MOL_TO_EV
+        )  # Convert to eV/Å for convergence criteria
+        return energy, grad.reshape(-1, 3)
+
+    def get_energies(self, batch: Data | Batch) -> torch.Tensor:
+        """Compute energies for a batch of conformers using MMFF94."""
+        self._ensure_single_conformer(batch)
+        mol = self._to_mol(batch)
+        mol = self._prepare_mol(mol)
+
+        ff = self._build_forcefield(mol)
+        return torch.tensor([ff.CalcEnergy()])

src/neural_optimiser/calculators/base.py

@@ -16,11 +16,9 @@ def __call__(self, batch: Data | Batch) -> tuple[torch.Tensor, torch.Tensor]:
         self.device = batch.pos.device
         return self._calculate(batch)
 
-    @abstractmethod
     def __repr__(self):
         ...
 
-    @abstractmethod
     def get_energies(self, batch: Data | Batch) -> torch.Tensor:
         """Get only energies from the calculator."""
         ...
@@ -30,7 +28,6 @@ def _calculate(self, batch: Data | Batch) -> tuple[torch.Tensor, torch.Tensor]:
         """Return (energies, forces) from the calculator."""
         ...
 
-    @abstractmethod
     def to_atomic_data(self, batch: Data | Batch) -> Batch:
         """Convert to AtomicData format compatible with ML model used."""
         ...

tests/calculators/test_mmff94.py

@@ -0,0 +1,129 @@
+import pytest
+import torch
+from ase.units import eV, kcal, mol
+from neural_optimiser.calculators import MMFF94Calculator
+from neural_optimiser.conformers import Conformer
+from rdkit import Chem
+from rdkit.Chem import AllChem, rdDetermineBonds
+
+KCAL_PER_MOL_TO_EV = kcal / mol / eV
+
+
+def test_MMFF94Calculator_calculate(mol):
+    """Compare energy and forces to RDKit's MMFF94 implementation."""
+    mp = AllChem.MMFFGetMoleculeProperties(mol)
+    ff = AllChem.MMFFGetMoleculeForceField(mol, mp)
+    _energy = ff.CalcEnergy()
+    _forces = torch.tensor(ff.CalcGrad()) * -KCAL_PER_MOL_TO_EV
+
+    data = Conformer.from_rdkit(mol)
+    calc = MMFF94Calculator()
+    energy, forces = calc._calculate(data)
+
+    assert torch.isclose(energy, torch.tensor(_energy))
+    assert torch.allclose(forces.flatten(), torch.Tensor(_forces))
+
+
+def test_repr_no_args():
+    """Test the __repr__ method with no arguments."""
+    calc = MMFF94Calculator()
+    assert repr(calc) == "MMFFCalculator()"
+
+
+def test_repr_with_args_only_affects_repr_not_behavior():
+    """Test the __repr__ method with arguments."""
+    calc = MMFF94Calculator(
+        MMFFGetMoleculeProperties={"mmffVariant": "MMFF94s"},
+        MMFFGetMoleculeForceField={"nonBondedThresh": 10.0},
+    )
+    s = repr(calc)
+    assert "mmffVariant=MMFF94s" in s
+    assert "nonBondedThresh=10.0" in s
+
+
+def test_single_conformer_energy_and_forces_shape(batch):
+    """Test that energy is scalar and forces have correct shape."""
+    calc = MMFF94Calculator()
+    energy, forces = calc._calculate(batch)
+
+    # energy is a scalar tensor
+    assert isinstance(energy, torch.Tensor)
+    assert energy.ndim == 1
+
+    # forces match [n_atoms, 3]
+    assert isinstance(forces, torch.Tensor)
+    assert forces.ndim == 2 and forces.shape[1] == 3
+    assert forces.shape[0] == batch.pos.shape[0]
+
+
+def test_get_energies_matches_calculate(batch):
+    """Test that get_energies matches the energy from _calculate."""
+    calc = MMFF94Calculator()
+    e_only = calc.get_energies(batch)
+    e_calc, _ = calc._calculate(batch)
+
+    assert torch.allclose(e_only, e_calc)
+
+
+def test_multi_conformer_raises(minimised_batch):
+    """Test that multi-conformer batches raise ValueError."""
+    calc = MMFF94Calculator()
+    with pytest.raises(ValueError):
+        calc._calculate(minimised_batch)
+    with pytest.raises(ValueError):
+        calc.get_energies(minimised_batch)
+
+
+def test_bond_determination_cached(monkeypatch, batch):
+    """Test that bond determination is cached based on molecule key."""
+    calls = {"n": 0}
+    orig = rdDetermineBonds.DetermineBonds
+
+    def wrapper(mol, *args, **kwargs):
+        calls["n"] += 1
+        return orig(mol, *args, **kwargs)
+
+    monkeypatch.setattr(rdDetermineBonds, "DetermineBonds", wrapper)
+
+    calc = MMFF94Calculator()
+
+    # First call should determine bonds once
+    _ = calc._calculate(batch)
+    assert calls["n"] == 1
+
+    # Second prepare on a new mol with same key should reuse cached bonds
+    mol2 = Chem.RWMol(batch.to_rdkit()[0])
+    _ = calc._prepare_mol(mol2)
+    assert calls["n"] == 1  # unchanged
+
+
+def test_smiles_property_controls_cache_key(monkeypatch, batch):
+    """Test that the 'smiles' property controls the caching of bond determination."""
+    calls = {"n": 0}
+    orig = rdDetermineBonds.DetermineBonds
+
+    def wrapper(mol, *args, **kwargs):
+        calls["n"] += 1
+        return orig(mol, *args, **kwargs)
+
+    monkeypatch.setattr(rdDetermineBonds, "DetermineBonds", wrapper)
+
+    calc = MMFF94Calculator()
+
+    # First molecule with a specific 'smiles' property
+    mol1 = Chem.RWMol(batch.to_rdkit()[0])
+    mol1.SetProp("smiles", "KEY")
+    _ = calc._prepare_mol(mol1)
+    assert calls["n"] == 1
+
+    # Same key => should not re-run DetermineBonds
+    mol2 = Chem.RWMol(batch.to_rdkit()[0])
+    mol2.SetProp("smiles", "KEY")
+    _ = calc._prepare_mol(mol2)
+    assert calls["n"] == 1
+
+    # Different key => should run DetermineBonds again
+    mol3 = Chem.RWMol(batch.to_rdkit()[0])
+    mol3.SetProp("smiles", "KEY2")
+    _ = calc._prepare_mol(mol3)
+    assert calls["n"] == 2

tests/optimise/test_bfgs.py

@@ -3,7 +3,7 @@
 import pytest
 import torch
 from neural_optimiser import test_dir
-from neural_optimiser.calculators import MACECalculator
+from neural_optimiser.calculators import MACECalculator, MMFF94Calculator
 from neural_optimiser.conformers import Conformer, ConformerBatch
 from neural_optimiser.optimisers import BFGS
 
@@ -115,6 +115,16 @@ def test_bfgs_integration(atoms, atoms2):
     assert converged is True
 
 
+def test_bfgs_integration2(atoms2):
+    """Test BFGS integration with MMFF94Calculator on CPU."""
+    batch = ConformerBatch.from_ase([atoms2])
+
+    optimiser = BFGS(steps=100, fmax=0.05, fexit=500.0)
+    optimiser.calculator = MMFF94Calculator()
+    converged = optimiser.run(batch)
+    assert converged is True
+
+
 def test_bfgs_integration_gpu(atoms, atoms2):
     """Test BFGS integration with MACECalculator on GPU."""
     pytest.importorskip("mace", reason="MACE not installed")