Add preprocessing of coordinates and switching function

[JMorado]

Now that EMLE is starting to be integrated into other simulation software (e.g., OpenMM), it is important to ensure that emle-engine contains all the necessary routines to reproduce the reference Sire implementation (except for Sire link-atom schemes).

To achieve this, this PR adds the following:

1. Coordinate preprocessing: This makes whole the QM region, re-images the MM atoms to their minimum image position with respect to the QM region centre, and applies a hard distance cutoff, zeroing the charges of MM atoms further than cutoff from the nearest QM atom.

2. Quintic switching function: This is applied to the MM charges to smooth the transition between the cutoff and non-cutoff regions.

Both features can be optionally enabled or disabled via the preprocess and use_switching_function flags in the forward EMLE model. By default, both are set to None, ensuring that EMLE continues working exactly as before.

cc @epretti

[lohedges]

Many thanks, @JMorado. I've tested locally and everything is working as expected. Ideally we should test forces as well as energies, since Sire handles the chain-rule contribution manually, whereas that would be done automatically via autograd in OpenMM-ML, e.g something like:

# Reference forces from OpenMM (kJ/mol/nm), QM force group only.
state = context.getState(getForces=True, groups={qm_force.getForceGroup()})
forces_ref = torch.tensor(
    state.getForces(asNumpy=True).value_in_unit(
        openmm.unit.kilojoule_per_mole / openmm.unit.nanometer
    ),
    dtype=dtype,
    device=device,
)
# Convert to Hartree/Å.
forces_ref = forces_ref / hartree_to_kj_mol / 10.0

# Model forces via autograd.
xyz_qm.requires_grad_(True)
xyz_mm_all.requires_grad_(True)
energy_test = model(
    atomic_numbers, charges_mm_all, xyz_
    cell, 0, None, True, use_switching_function,
)
grad_qm, grad_mm = torch.autograd.grad(ez_mm_all))
forces_test = -torch.cat([grad_qm, grad_

assert torch.allclose(forces_ref, forces)

(Although this is more a test that Sire is correct, not emle-engine.)

[JMorado]

Thanks @lohedges.

Yes, that makes sense to me. Even if this is more of a test to make sure Sire is correct, I think it's still important to ensure that both implementations are on par, so I'll add your code snippet.

[lohedges]

Great, thanks for adding that!

[lohedges]

Just realised that this went to main, not devel. This is fine, since it's a API backwards compatible. Could you possibly create another PR with the same changes to devel too? I can always cherry-pick, but would be good to have a paper trail.

[JMorado]

Oops, my bad! Yes, sure. I'll create another PR right away.