D3 dispersion corrections with MaceModel

[yjw0510]

Dear TorchSim users and developers,

I'm trying to port an ASE MD script that uses mace_mp(dispersion=True) over to TorchSim.

The problem is that to use MACE with TorchSim, you need return_raw_model=True:

mace = mace_mp(model="large", return_raw_model=True)
mace_model = MaceModel(model=mace, device=device)

But looking at the mace_mp source, when return_raw_model=True it just does torch.load() and returns immediately. The D3(BJ) dispersion part (which wraps things in an ASE SumCalculator) is never reached. So the raw model you get back has no dispersion at all.

Then on the TorchSim side, MaceModel takes this raw nn.Module and wraps it as-is. I couldn't find anything like a model summing/composition API that would let you bolt on a separate D3 calculation.

Is there a way to do this that I'm not seeing, or is dispersion correction just not supported yet?

Thanks and Best Regards,
Ji Woong

[CompRhys]

This would need to be implemented inside the model. We do not have an equivalent to a SumCalculator yet it is something we could consider adding. We're moving to an external model posture and a torchsim interface has been added to acesuit/mace on github so potentially a solution already exists for you there.

[CompRhys]

A sum calculator and D3 model have now been added to torchsim. Please take a look and let us know if you have any issues.