Conversation
|
Hmm, confused by the conflicts here. I'll take a look from the command-line to see if I can resolve them. |
|
@lohedges, this would add all the commits from main to |
|
Ah, I see. Yes, sorry if I wasn't clear. Currently we just want the new commits on top of |
Adds DeePMDEMLE, a torch.nn.Module composite mirroring MACEEMLE that loads a DeePMD-kit v3 PyTorch-backend .pth (TorchScript) model via torch.jit.load and embeds it as a submodule alongside the EMLE embedding head, returning (E_vac, E_static, E_induced) in Hartree. Atomic numbers are mapped to DeePMD-internal type indices via a 119-element buffer built once from the model's get_type_map(). TensorFlow .pb models are rejected with a clear error pointing at the runtime DeePMD backend. Tests script the composite under TorchScript and exercise unbatched and batched forward passes with autograd through the QM and MM coordinates. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The .to/.cpu/.cuda/.double/.float overrides only forwarded the call to the EMLE and DeePMD submodules, so the directly-registered buffers (_atomic_numbers, _z_to_type) stayed on the original device after .cuda()/.cpu() — breaking forward at self._z_to_type[atomic_numbers]. Add super() calls so torch.nn.Module moves the buffers, and extend .to() to track dtype kwargs/args alongside device. Also fix two dtype leaks in forward: DeePMD unconditionally casts its energy_redu output to float64, so cast E_vac back to self._dtype, and build the empty-MM zeros row from self._dtype rather than xyz_qm.dtype to avoid silent dtype promotion in torch.stack when those differ. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The constructor now also accepts a list/tuple of model paths, in which case the wrapper loads all members into a ModuleList. The first member remains the primary model used for the returned in-vacuo energy; on forward, every member is run and per-model energies and DeePMD-native forces are stored on _E_vac_qbc and _grads_qbc in the shape consumed by EMLECalculator's deviation computation. grads = -force matches MACEEMLE's autograd-of-energy convention so the two backends are interchangeable downstream. Ensemble members must share a type_map; otherwise the same atype tensor would be silently misinterpreted by the secondaries, so the constructor cross-checks and raises. The .to/.cpu/.cuda/.double/.float overrides iterate over the ModuleList, and self._deepmd is rebound after each move so it always references the same underlying module as self._deepmd_models[0]. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
EMLECompiler now accepts backend="deepmd" and a deepmd_model kwarg, and the emle-compile CLI exposes --deepmd-model. The compile path always wraps a single in-vacuo model; ensemble (QbC) usage stays at the DeePMDEMLE constructor for runtime EMLECalculator. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Is it possible to rebase this and force push? There should really only be one new commit here, not 59. Alternatively, create a new branch based off of the current |
|
Cherry-picking got a bit too complicated, and I'm worried I might have accidentally changed some important code while resolving conflicts. I'll just open a clean PR instead. |
JMorado
changed the title
|
No problem. Sounds like the easiest solution 👍 |
3 participants
Clone of #88 for
devel, keeping the same changes for a clear paper trail.