From 6d5274f7048bf8ab61a0a85fd05fac6601f5f00f Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 28 Sep 2026 20:11:41 +0100 Subject: [PATCH] Backport fix from PR #492. [ci skip] --- src/sire/qm/_utils.py | 20 ++- tests/qm/test_qm.py | 172 +++++++++++++++++++++++++ wrapper/Convert/SireOpenMM/pyqm.cpp | 115 +++++++++++++---- wrapper/Convert/SireOpenMM/torchqm.cpp | 128 ++++++++++++++---- 4 files changed, 382 insertions(+), 53 deletions(-) diff --git a/src/sire/qm/_utils.py b/src/sire/qm/_utils.py index a3529db78..ddd7e8b79 100644 --- a/src/sire/qm/_utils.py +++ b/src/sire/qm/_utils.py @@ -409,6 +409,24 @@ def _get_link_atoms(mols, qm_mol_to_atoms, map): mm_bonds.append(bond_idx) mm2_atoms[mm1_idx] = mm_bonds + # The charge shift method can't handle MM1 atoms that are bonded to + # each other, or MM2 atoms that are shared between MM1 atoms. + seen_mm2 = set() + for mm1_idx, mm_bonds in mm2_atoms.items(): + for mm2_idx in mm_bonds: + abs_idx = mols.atoms().find(qm_mol.atoms()[mm2_idx]) + if mm2_idx in mm2_atoms: + raise Exception( + f"MM atom (index {abs_idx}) is bonded to the QM region and " + "to another MM atom that is bonded to the QM region!" + ) + if mm2_idx in seen_mm2: + raise Exception( + f"MM atom (index {abs_idx}) is bonded to more than one " + "MM atom that is bonded to the QM region!" + ) + seen_mm2.add(mm2_idx) + # Convert MM1 to QM atom dictionary to absolute indices. mm1_to_qm_local = {} for k, v in mm1_atoms.items(): @@ -505,7 +523,7 @@ def _get_link_atoms(mols, qm_mol_to_atoms, map): if not qm_link_bond_found: # First find a standard C-H bond length. carbon = _Mol.Element("C") - bonds = qm_mol[f"element C"].bonds() + bonds = qm_mol["element C"].bonds() ch_bond_length = None for bond in bonds: atom0 = bond.atom0() diff --git a/tests/qm/test_qm.py b/tests/qm/test_qm.py index 0d5f3d104..22f61747c 100644 --- a/tests/qm/test_qm.py +++ b/tests/qm/test_qm.py @@ -187,6 +187,31 @@ def test_link_atoms_non_carbon(): ) +@pytest.mark.parametrize( + "selection, match", + [ + # MM1 atoms 4 and 6 are bonded to each other. + ("atomidx 0:4 or atomidx 8:22", "to another MM atom"), + # MM1 atoms 4 and 8 share MM2 atom 6. + ("atomidx 0:4 or atomidx 14:22", "more than one MM atom"), + ], +) +def test_link_atoms_unsupported(ala_mols, selection, match): + """ + Make sure that link atom layouts the charge shift method can't handle + are rejected. + """ + + from sire.base import create_map as _create_map + from sire.qm._utils import _create_qm_mol_to_atoms, _get_link_atoms + + qm_atoms = ala_mols[0][selection].atoms() + qm_mol_to_atoms = _create_qm_mol_to_atoms(qm_atoms) + + with pytest.raises(Exception, match=match): + _get_link_atoms(ala_mols, qm_mol_to_atoms, _create_map({})) + + def test_charge_redistribution(): """ Make sure that charge redistribution works correctly. @@ -453,6 +478,153 @@ def callback(numbers_qm, charges_mm, xyz_qm, xyz_mm, cell=None, idx_mm=None): assert nrg == 42 +def _check_link_atom_forces(qm_mols, engine, atol): + """ + Compare the forces on the atoms around the residue 1 link atoms of the + alanine dipeptide to central finite differences of the total energy. + """ + + import openmm + + d = qm_mols[0].dynamics( + timestep="1fs", + constraint="none", + qm_engine=engine, + cutoff_type="pme", + cutoff="7.5 A", + platform="reference", + ) + + context = d.context() + + nm = openmm.unit.nanometer + kj = openmm.unit.kilojoule_per_mole + + state = context.getState(getPositions=True, getForces=True) + positions = state.getPositions(asNumpy=True).value_in_unit(nm) + forces = state.getForces(asNumpy=True).value_in_unit(kj / nm) + + def energy(pos): + context.setPositions(pos * nm) + return context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(kj) + + # MM1 atoms 4 and 16 are bonded to QM1 atoms 6 and 14, with MM2 atoms + # 1, 5 and 17, 18. Atom 10 is a QM control. + delta = 1e-4 + for idx in [1, 4, 5, 6, 10, 14, 16, 17, 18]: + fd_force = np.zeros(3) + for k in range(3): + pos = positions.copy() + pos[idx, k] += delta + e_plus = energy(pos) + pos[idx, k] -= 2 * delta + e_minus = energy(pos) + fd_force[k] = -(e_plus - e_minus) / (2 * delta) + + assert np.allclose(forces[idx], fd_force, atol=atol) + + +@pytest.mark.parametrize("mechanical_embedding", [False, True]) +def test_link_atom_forces(ala_mols, mechanical_embedding): + """ + Make sure that the forces on the QM1, MM1 and MM2 atoms around each link + atom are the negative gradient of the QM energy. + """ + + # Energy coefficients (kJ/mol/A) for the two link atoms, which are the + # final rows of xyz_qm. + link_coeffs = np.array([[50.0, -30.0, 20.0], [-40.0, 25.0, 60.0]]) + + # Energy coefficients (kJ/mol/A) for each pair of virtual point charges, + # which are the final rows of xyz_mm. + virtual_coeffs = np.array([[10.0, -5.0, 8.0], [-7.0, 12.0, 3.0]]) + + # The number of QM rows passed to the callback. + num_rows = [] + + # A callback returning an energy that is linear in the QM and virtual + # point charge positions. + def callback(numbers_qm, charges_mm, xyz_qm, xyz_mm, cell=None, idx_mm=None): + num_rows.append(len(xyz_qm)) + num_qm = len(xyz_qm) - len(link_coeffs) + coeffs_qm = np.vstack( + [0.1 * np.outer(np.arange(1, num_qm + 1), [1.0, 2.0, 3.0]), link_coeffs] + ) + num_pairs = (len(xyz_mm) - len(idx_mm)) // 2 + coeffs_mm = np.vstack( + [np.zeros((len(idx_mm), 3)), np.tile(virtual_coeffs, (num_pairs, 1))] + ) + energy = float( + np.sum(coeffs_qm * np.array(xyz_qm)) + + np.sum(coeffs_mm * np.array(xyz_mm).reshape(-1, 3)) + ) + return (energy, (-10.0 * coeffs_qm).tolist(), (-10.0 * coeffs_mm).tolist()) + + # Create a local copy of the test system. + mols = ala_mols.clone() + + qm_mols, engine = sr.qm.create_engine( + mols, + mols[0]["residx 1"], + callback, + mechanical_embedding=mechanical_embedding, + ) + + _check_link_atom_forces(qm_mols, engine, atol=0.1) + + # Make sure that both link atoms were added to the QM region. + assert num_rows[-1] == mols[0]["residx 1"].num_atoms() + 2 + + +@pytest.mark.skipif(not has_emle, reason="emle-engine is not installed") +def test_link_atom_forces_torch(ala_mols, tmp_path, monkeypatch): + """ + Make sure that the forces on the QM1, MM1 and MM2 atoms around each link + atom are the negative gradient of the energy of a TorchScript QM model. + """ + + import torch + + # The engine saves the TorchScript module to the working directory. + monkeypatch.chdir(tmp_path) + + # A model returning an energy (Hartree) that is linear in the QM and MM + # positions, with distinct coefficients for the two link atoms, which + # are the final rows of xyz_qm. + class LinearModel(torch.nn.Module): + _is_emle = True + + def forward( + self, + atomic_numbers: torch.Tensor, + charges_mm: torch.Tensor, + xyz_qm: torch.Tensor, + xyz_mm: torch.Tensor, + cell: torch.Tensor, + ) -> torch.Tensor: + assert xyz_qm.shape[0] == 12 + coeffs = 1e-4 * torch.outer( + torch.arange(1, 13, dtype=torch.float64), + torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64), + ) + coeffs[10] = torch.tensor([5e-3, -3e-3, 2e-3], dtype=torch.float64) + coeffs[11] = torch.tensor([-4e-3, 2.5e-3, 6e-3], dtype=torch.float64) + mm_coeffs = torch.tensor([1e-3, -5e-4, 8e-4], dtype=torch.float64) + return (coeffs * xyz_qm.double()).sum() + ( + mm_coeffs * xyz_mm.double() + ).sum() + + # Create a local copy of the test system. + mols = ala_mols.clone() + + qm_mols, engine = sr.qm.emle( + mols, mols[0]["residx 1"], LinearModel(), switch_width=0.0 + ) + + # The model uses single precision positions, so allow for rounding. + _check_link_atom_forces(qm_mols, engine, atol=1.0) + + def test_qmff_was_force_changed(ala_mols): """ Verify that wasForceChanged("qmff") correctly tracks whether the qmff diff --git a/wrapper/Convert/SireOpenMM/pyqm.cpp b/wrapper/Convert/SireOpenMM/pyqm.cpp index 24e9ede63..b7f5b1631 100644 --- a/wrapper/Convert/SireOpenMM/pyqm.cpp +++ b/wrapper/Convert/SireOpenMM/pyqm.cpp @@ -49,7 +49,7 @@ using namespace SireStream; using namespace SireVol; // The delta used to place virtual point charges either side of the MM2 -// atoms, in nanometers. +// atoms, in Angstrom. static const double VIRTUAL_PC_DELTA = 0.01; class GILLock @@ -722,6 +722,53 @@ double PyQMForceImpl::computeForce( QVector nearest_qm_vecs; QVector nearest_qm_atom_idxs; + // The QM1 and MM1 indices and bond scale factor for each link atom. + QVector link_qm1_idxs; + QVector link_mm1_idxs; + QVector link_scales; + + // The MM1 and MM2 indices, MM1-MM2 unit vector and bond length for each + // pair of virtual point charges. + QVector virtual_mm1_idxs; + QVector virtual_mm2_idxs; + QVector virtual_normals; + QVector virtual_lengths; + + // Add a link atom to the QM region for each QM-MM1 bond. + for (const auto &idx : mm1_to_mm2.keys()) + { + // Get the QM atom to which the current MM atom is bonded. + const auto qm_idx = mm1_to_qm[idx]; + + // Store the MM1 position in Sire Vector format, along with the + // position of the QM atom to which it is bonded. + Vector mm1_vec(10 * positions[idx][0], 10 * positions[idx][1], 10 * positions[idx][2]); + Vector qm_vec(10 * positions[qm_idx][0], 10 * positions[qm_idx][1], 10 * positions[qm_idx][2]); + + // Work out the minimum image positions with respect to the reference position. + mm1_vec = space.getMinimumImage(mm1_vec, center); + qm_vec = space.getMinimumImage(qm_vec, center); + + // Work out the position of the link atom. Here we use a bond length + // scale factor taken from the MM bond potential, i.e. R0(QM-L) / R0(QM-MM1), + // where R0(QM-L) is the equilibrium bond length for the QM and link (L) + // elements, and R0(QM-MM1) is the equilibrium bond length for the QM + // and MM1 elements. + const auto link_vec = qm_vec + bond_scale_factors[idx] * (mm1_vec - qm_vec); + + // Add to the QM positions. + xyz_qm.append(QVector({link_vec[0], link_vec[1], link_vec[2]})); + + // Store the link atom info so that its force can be split between + // the QM1 and MM1 atoms. + link_qm1_idxs.append(qm_idx); + link_mm1_idxs.append(idx); + link_scales.append(bond_scale_factors[idx]); + + // Append a hydrogen element to the numbers vector. + numbers.append(1); + } + // If we are using electrostatic embedding, the work out the MM point charges and // build the neighbour list. if (not this->owner.getIsMechanical()) @@ -865,33 +912,11 @@ double PyQMForceImpl::computeForce( // See: https://www.ks.uiuc.edu/Research/qmmm for (const auto &idx : mm1_to_mm2.keys()) { - // Get the QM atom to which the current MM atom is bonded. - const auto qm_idx = mm1_to_qm[idx]; - - // Store the MM1 position in Sire Vector format, along with the - // position of the QM atom to which it is bonded. + // Store the MM1 position in Sire Vector format. Vector mm1_vec(10 * positions[idx][0], 10 * positions[idx][1], 10 * positions[idx][2]); - Vector qm_vec(10 * positions[qm_idx][0], 10 * positions[qm_idx][1], 10 * positions[qm_idx][2]); - // Work out the minimum image positions with respect to the reference position. + // Work out the minimum image position with respect to the reference position. mm1_vec = space.getMinimumImage(mm1_vec, center); - qm_vec = space.getMinimumImage(qm_vec, center); - - // Work out the position of the link atom. Here we use a bond length - // scale factor taken from the MM bond potential, i.e. R0(QM-L) / R0(QM-MM1), - // where R0(QM-L) is the equilibrium bond length for the QM and link (L) - // elements, and R0(QM-MM1) is the equilibrium bond length for the QM - // and MM1 elements. - const auto link_vec = qm_vec + bond_scale_factors[idx] * (mm1_vec - qm_vec); - - // Add to the QM positions. - xyz_qm.append(QVector({link_vec[0], link_vec[1], link_vec[2]})); - - // Add the MM1 index to the QM atoms vector. - qm_atoms.append(qm_idx); - - // Append a hydrogen element to the numbers vector. - numbers.append(1); // Store the number of MM2 atoms. const auto num_mm2 = mm1_to_mm2[idx].size(); @@ -924,6 +949,13 @@ double PyQMForceImpl::computeForce( // Compute the normal vector from the MM1 to MM2 atom. const auto normal = (mm2_vec - mm1_vec).normalise(); + // Store the info needed to project the virtual point charge + // forces onto the MM1 and MM2 atoms. + virtual_mm1_idxs.append(idx); + virtual_mm2_idxs.append(mm2_idx); + virtual_normals.append(normal); + virtual_lengths.append((mm2_vec - mm1_vec).length()); + // Positive direction. (Away from MM1 atom.) auto xyz = mm2_vec + VIRTUAL_PC_DELTA * normal; xyz_virtual.append(QVector({xyz[0], xyz[1], xyz[2]})); @@ -1010,6 +1042,19 @@ double PyQMForceImpl::computeForce( forces[idx] = lambda * omm_force; } + // Split the link atom forces between the QM1 and MM1 atoms, since the + // link atom position is L = QM1 + g * (MM1 - QM1). + for (int j = 0; j < link_qm1_idxs.size(); j++) + { + const auto i = qm_atoms.size() + j; + + OpenMM::Vec3 omm_force(forces_qm[i][0], forces_qm[i][1], forces_qm[i][2]); + + const auto g = link_scales[j]; + forces[link_qm1_idxs[j]] += lambda * (1.0 - g) * omm_force; + forces[link_mm1_idxs[j]] += lambda * g * omm_force; + } + // Now the MM atoms. for (int i = 0; i < num_mm; i++) { @@ -1048,6 +1093,26 @@ double PyQMForceImpl::computeForce( } } + // Project the virtual point charge forces onto the MM1 and MM2 atoms, since + // their positions are MM2 +/- delta * n, where n is the MM1-MM2 unit vector. + for (int j = 0; j < virtual_mm1_idxs.size(); j++) + { + const auto i = num_mm + 2 * j; + + const Vector f_plus(forces_mm[i][0], forces_mm[i][1], forces_mm[i][2]); + const Vector f_minus(forces_mm[i + 1][0], forces_mm[i + 1][1], forces_mm[i + 1][2]); + + // The component from the change in direction of n. + const auto &n = virtual_normals[j]; + const auto f_diff = f_plus - f_minus; + const auto f_rot = (VIRTUAL_PC_DELTA / virtual_lengths[j]) * + (f_diff - Vector::dot(f_diff, n) * n); + + const auto f_mm2 = f_plus + f_minus + f_rot; + forces[virtual_mm2_idxs[j]] += lambda * OpenMM::Vec3(f_mm2[0], f_mm2[1], f_mm2[2]); + forces[virtual_mm1_idxs[j]] -= lambda * OpenMM::Vec3(f_rot[0], f_rot[1], f_rot[2]); + } + // Update the step count. this->step_count++; diff --git a/wrapper/Convert/SireOpenMM/torchqm.cpp b/wrapper/Convert/SireOpenMM/torchqm.cpp index 40c56acb2..b8ceb5bf2 100644 --- a/wrapper/Convert/SireOpenMM/torchqm.cpp +++ b/wrapper/Convert/SireOpenMM/torchqm.cpp @@ -49,7 +49,7 @@ using namespace SireStream; using namespace SireVol; // The delta used to place virtual point charges either side of the MM2 -// atoms, in nanometers. +// atoms, in Angstrom. static const double VIRTUAL_PC_DELTA = 0.01; // Conversion factor from Hartree to kJ/mol. @@ -539,6 +539,55 @@ double TorchQMForceImpl::computeForce( QVector nearest_qm_vecs; QVector nearest_qm_atom_idxs; + // The QM1 and MM1 indices and bond scale factor for each link atom. + QVector link_qm1_idxs; + QVector link_mm1_idxs; + QVector link_scales; + + // The MM1 and MM2 indices, MM1-MM2 unit vector and bond length for each + // pair of virtual point charges. + QVector virtual_mm1_idxs; + QVector virtual_mm2_idxs; + QVector virtual_normals; + QVector virtual_lengths; + + // Add a link atom to the QM region for each QM-MM1 bond. + for (const auto &idx : mm1_to_mm2.keys()) + { + // Get the QM atom to which the current MM atom is bonded. + const auto qm_idx = mm1_to_qm[idx]; + + // Store the MM1 position in Sire Vector format, along with the + // position of the QM atom to which it is bonded. + Vector mm1_vec(10 * positions[idx][0], 10 * positions[idx][1], 10 * positions[idx][2]); + Vector qm_vec(10 * positions[qm_idx][0], 10 * positions[qm_idx][1], 10 * positions[qm_idx][2]); + + // Work out the minimum image positions with respect to the reference position. + mm1_vec = space.getMinimumImage(mm1_vec, center); + qm_vec = space.getMinimumImage(qm_vec, center); + + // Work out the position of the link atom. Here we use a bond length + // scale factor taken from the MM bond potential, i.e. R0(QM-L) / R0(QM-MM1), + // where R0(QM-L) is the equilibrium bond length for the QM and link (L) + // elements, and R0(QM-MM1) is the equilibrium bond length for the QM + // and MM1 elements. + const auto link_vec = qm_vec + bond_scale_factors[idx] * (mm1_vec - qm_vec); + + // Add to the QM positions. + xyz_qm.push_back(link_vec[0]); + xyz_qm.push_back(link_vec[1]); + xyz_qm.push_back(link_vec[2]); + + // Store the link atom info so that its force can be split between + // the QM1 and MM1 atoms. + link_qm1_idxs.append(qm_idx); + link_mm1_idxs.append(idx); + link_scales.append(bond_scale_factors[idx]); + + // Append a hydrogen element to the numbers vector. + numbers.append(1); + } + // If we are using electrostatic embedding, the work out the MM point charges and // build the neighbour list. if (not this->owner.getIsMechanical()) @@ -686,35 +735,11 @@ double TorchQMForceImpl::computeForce( // See: https://www.ks.uiuc.edu/Research/qmmm for (const auto &idx : mm1_to_mm2.keys()) { - // Get the QM atom to which the current MM atom is bonded. - const auto qm_idx = mm1_to_qm[idx]; - - // Store the MM1 position in Sire Vector format, along with the - // position of the QM atom to which it is bonded. + // Store the MM1 position in Sire Vector format. Vector mm1_vec(10 * positions[idx][0], 10 * positions[idx][1], 10 * positions[idx][2]); - Vector qm_vec(10 * positions[qm_idx][0], 10 * positions[qm_idx][1], 10 * positions[qm_idx][2]); - // Work out the minimum image positions with respect to the reference position. + // Work out the minimum image position with respect to the reference position. mm1_vec = space.getMinimumImage(mm1_vec, center); - qm_vec = space.getMinimumImage(qm_vec, center); - - // Work out the position of the link atom. Here we use a bond length - // scale factor taken from the MM bond potential, i.e. R0(QM-L) / R0(QM-MM1), - // where R0(QM-L) is the equilibrium bond length for the QM and link (L) - // elements, and R0(QM-MM1) is the equilibrium bond length for the QM - // and MM1 elements. - const auto link_vec = qm_vec + bond_scale_factors[idx] * (mm1_vec - qm_vec); - - // Add to the QM positions. - xyz_qm.push_back(link_vec[0]); - xyz_qm.push_back(link_vec[1]); - xyz_qm.push_back(link_vec[2]); - - // Add the MM1 index to the QM atoms vector. - qm_atoms.append(qm_idx); - - // Append a hydrogen element to the numbers vector. - numbers.append(1); // Store the number of MM2 atoms. const auto num_mm2 = mm1_to_mm2[idx].size(); @@ -749,6 +774,13 @@ double TorchQMForceImpl::computeForce( // Compute the normal vector from the MM1 to MM2 atom. const auto normal = (mm2_vec - mm1_vec).normalise(); + // Store the info needed to project the virtual point charge + // forces onto the MM1 and MM2 atoms. + virtual_mm1_idxs.append(idx); + virtual_mm2_idxs.append(mm2_idx); + virtual_normals.append(normal); + virtual_lengths.append((mm2_vec - mm1_vec).length()); + // Positive direction. (Away from MM1 atom.) auto xyz = mm2_vec + VIRTUAL_PC_DELTA * normal; xyz_virtual.push_back(xyz[0]); @@ -926,6 +958,22 @@ double TorchQMForceImpl::computeForce( forces[idx] = lambda * omm_force; } + // Split the link atom forces between the QM1 and MM1 atoms, since the + // link atom position is L = QM1 + g * (MM1 - QM1). + for (int j = 0; j < link_qm1_idxs.size(); j++) + { + const auto i = qm_atoms.size() + j; + + OpenMM::Vec3 omm_force( + forces_qm_flat[3 * i], + forces_qm_flat[3 * i + 1], + forces_qm_flat[3 * i + 2]); + + const auto g = link_scales[j]; + forces[link_qm1_idxs[j]] += lambda * (1.0 - g) * omm_force; + forces[link_mm1_idxs[j]] += lambda * g * omm_force; + } + // Now the MM atoms. for (int i = 0; i < num_mm; i++) { @@ -965,6 +1013,32 @@ double TorchQMForceImpl::computeForce( } } + // Project the virtual point charge forces onto the MM1 and MM2 atoms, since + // their positions are MM2 +/- delta * n, where n is the MM1-MM2 unit vector. + for (int j = 0; j < virtual_mm1_idxs.size(); j++) + { + const auto i = num_mm + 2 * j; + + const Vector f_plus( + forces_mm_flat[3 * i], + forces_mm_flat[3 * i + 1], + forces_mm_flat[3 * i + 2]); + const Vector f_minus( + forces_mm_flat[3 * i + 3], + forces_mm_flat[3 * i + 4], + forces_mm_flat[3 * i + 5]); + + // The component from the change in direction of n. + const auto &n = virtual_normals[j]; + const auto f_diff = f_plus - f_minus; + const auto f_rot = (VIRTUAL_PC_DELTA / virtual_lengths[j]) * + (f_diff - Vector::dot(f_diff, n) * n); + + const auto f_mm2 = f_plus + f_minus + f_rot; + forces[virtual_mm2_idxs[j]] += lambda * OpenMM::Vec3(f_mm2[0], f_mm2[1], f_mm2[2]); + forces[virtual_mm1_idxs[j]] -= lambda * OpenMM::Vec3(f_rot[0], f_rot[1], f_rot[2]); + } + // Update the step count. this->step_count++;