CederGroupHub / chgnet

Pretrained universal neural network potential for charge-informed atomistic modeling https://chgnet.lbl.gov
https://doi.org/10.1038/s42256-023-00716-3
Other
226 stars 61 forks source link

[Feature Request]: Add dispersion to ASE calculator #188

Closed ajhoffman1229 closed 3 days ago

ajhoffman1229 commented 3 weeks ago

Email (Optional)

ajhoff29@mit.edu

Problem

I would like to be able to add a dispersion/van der Waals correction to the pre-trained CHGNet model. A similar option is available in the MACE-MP-0 model here.

Proposed Solution

I believe that an option can be added to the CHGNetCalculator for a dispersion correction similar to the option in MACE-MP-0. Their implementation requires installing torch-dftd, and creates an instance of an ASE SumCalculator that adds the energies, forces, and stresses of each contribution together.

Alternatives

If this feature cannot be added, it would be great if a note could be added to the README pointing users to the torch-dftd implementation in the event that they want to add a dispersion correction to CHGNet independent of the repo. Currently, the python packages for DFT-D3 and DFT-D4 dispersion models from the Grimme group are not compatible with torch, so the conventional method of adding the dispersion does not work with CHGNet.

Here is some example code where this works with CHGNet and the torch-dftd implementation:

from ase.build import fcc111
from ase.calculators.mixing import SumCalculator
from chgnet.model.dynamics import CHGNetCalculator
from chgnet.model.model import CHGNet
from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator

# pre-trained chgnet model
chgpotential = CHGNet.load()
chgcalc = CHGNetCalculator(model=chgpotential)

# Calculate reference values:
atoms = fcc111("Cu", (2, 2, 1), vacuum=10.0)
atoms[0].x += 0.2

chgd3 = SumCalculator([chgcalc, TorchDFTD3Calculator()])
atoms.calc = chgd3
stress = atoms.get_stress()
print(chgd3.name)
print(atoms.get_chemical_formula(), stress.shape, stress)

Code of Conduct

Andrew-S-Rosen commented 3 weeks ago

Currently, the python packages for DFT-D3 and DFT-D4 dispersion models from the Grimme group are not compatible with torch, so the conventional method of adding the dispersion does not work with CHGNet.

See https://github.com/dftd3/tad-dftd3 and https://github.com/dftd4/tad-dftd4.

ajhoffman1229 commented 3 weeks ago

Currently, the python packages for DFT-D3 and DFT-D4 dispersion models from the Grimme group are not compatible with torch, so the conventional method of adding the dispersion does not work with CHGNet.

See https://github.com/dftd3/tad-dftd3 and https://github.com/dftd4/tad-dftd4.

Thanks, I am aware of these repositories, but I don't think they yet have an implementation with ASE and can't be simply plugged into an ASE SumCalculator instance. I'd be willing to help develop something for them elsewhere (such as within ASE or adding the D4 implementation to the torch-dftd repo), but in the meantime, it would be great to fold in the existing torch-dftd implementations with CHGNet or create a tutorial showing users how to use them on their own.

Andrew-S-Rosen commented 3 weeks ago

Makes sense! Definitely in the meantime, torch-dftd seems like the way to go here.

BowenD-UCB commented 3 weeks ago

@ajhoffman1229 Can you please give some tests on the DFTD and SumCalculator and share a notebook example or open a PR? Happy to help and merge if the combined calculator works stably on your side.

If the above code works smoothly with no compatibility issues, the easiest thing is to add a DFT-D3.ipynb under chgnet/examples

ajhoffman1229 commented 3 weeks ago

Yes, I can make a notebook and submit a PR so it can be added as an example to the repo. The code above worked when I tested it locally. I will try to polish it later this week and submit for review.

ajhoffman1229 commented 2 weeks ago

Hi folks, I am not sure what I did but the code that I wrote earlier this week now seems to be working now and I can compute energies with a CHGNet+D4 SumCalculator instance in ASE, in addition to the torch-dftd D3 calculator. There does not appear to be the incompatibility I was running into earlier, and I apologize for any wild goose chases opening this issue might have caused. I made a notebook tutorial showing how to add a dispersion correction, which is in #192.

Andrew-S-Rosen commented 2 weeks ago

That's wonderful! Thanks for the update!