Backport fix from PR #492

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()

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

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<Vector> nearest_qm_vecs;
     QVector<int> nearest_qm_atom_idxs;
 
+    // The QM1 and MM1 indices and bond scale factor for each link atom.
+    QVector<int> link_qm1_idxs;
+    QVector<int> link_mm1_idxs;
+    QVector<double> link_scales;
+
+    // The MM1 and MM2 indices, MM1-MM2 unit vector and bond length for each
+    // pair of virtual point charges.
+    QVector<int> virtual_mm1_idxs;
+    QVector<int> virtual_mm2_idxs;
+    QVector<Vector> virtual_normals;
+    QVector<double> 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<double>({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<double>({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<double>({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++;