Allow custom mace model by specifying "model" in calculator kwargs"

[orionarcher]

Currently it's only possible to use the mace_mp model. This PR allows users to use a custom model by making the "model" kwarg a path to the model. If model is not a Path, then mace_mp is used, thus this change should be backwards compatible.

EDIT: More context:

mace==0.3.3, the version currently used by atomate2, allows custom mace models in mace_mp. So the current code does not need any changes to use custom models.

mace==0.3.7, the most recent version does not allow custom models in mace_mp, which is a breaking change. This PR will be needed when atomate2 bumps the mace version. When that happens, many tests will likely need fixing.

[orionarcher]

More context:

mace 0.3.3, the version currently used by atomate2, allows custom mace models in mace_mp. So the current code does not need any changes to use custom models.

mace 0.3.7, the most recent version does not allow custom models in mace_mp, which is a breaking change. This PR will be needed when atomate2 bumps the mace version. When that happens, many tests will likely need fixing.

@janosh @utf, how does atomate2 intend to evolve the mace dependency over time? They introduced some breaking changes but also useful new features.

[janosh]

@janosh @utf, how does atomate2 intend to evolve the mace dependency over time? They introduced some breaking changes but also useful new features.

i think the latest version of atomate2 should strive to be compatible with the latest version of mace-torch, even if that means adapting to breaking changes

[JaGeo]

@orionarcher I will remove the orb stuff here and then merge. There is a separate pull request for the orb part in #1031

[JaGeo]

Oh, there seems to be something wrong with the kwargs. Will check later and then merge.

[JaGeo]

@esoteric-ephemera Could you shortly have a look at this failing test? To me, it looks like the symmetry does not change during relaxation with the latest mace version. At least, within the ASERelaxer, the correct "relax_cell" and "fix_symmetry" are set.

[esoteric-ephemera]

@JaGeo sure! Will look at this soon

[esoteric-ephemera]

Think we need to change this to:

device = kwargs.get("device") or "cpu"
calculator =  MACECalculator(
    model_paths=model_path,
    device = device,
    **kwargs
)

[JaGeo]

Thanks. One needs to remove the "device" key from the kwargs. Otherwise, it will complain that it is set twice. Will do that.

[esoteric-ephemera]

OK so if the symmetry of a cell is fixed, we certainly should assert that the initial and final symmetries of the cell are the same

But if we don't fix the symmetry, there's a distinct possibility that the symmetry of the cell remains the same. It could change, but unless the force convergence criterion is made tighter, there's a decent chance the structure barely changes

For example, with the previous setting of fmax = 0.04 eV/Å, the atoms don't move much. With that tightened to 0.01 eV/Å, the symmetry of the cell changes during relaxation

Also worth noting that it seems like ASE's FixSymmetry isn't holding for more stringent force convergence criteria. I pushed fmax to 0.001 eV/Å and the cell symmetry wasn't preserved with fix_symmetry = True

[JaGeo]

@esoteric-ephemera so, you are okay with simply adapting the test, right?

[codecov]

Codecov Report

Attention: Patch coverage is 83.33333% with 2 lines in your changes missing coverage. Please review.

Project coverage is 66.58%. Comparing base (42bc7b8) to head (1f5e138).

Files with missing lines	Patch %	Lines
src/atomate2/forcefields/utils.py	83.33%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##            main    #1017       +/-   ##
==========================================
+ Coverage   4.32%   66.58%   +62.25%     
==========================================
  Files        178      178               
  Lines      12911    12921       +10     
  Branches    1278     1280        +2     
==========================================
+ Hits         559     8604     +8045     
+ Misses     12321     3838     -8483     
- Partials      31      479      +448     

Files with missing lines	Coverage Δ
src/atomate2/forcefields/utils.py	68.33% <83.33%> (+68.33%)	⬆️

... and 150 files with indirect coverage changes

[JaGeo]

@esoteric-ephemera @orionarcher Okay, I figured it out. mace_mp has a default floating type precision of float32. If we set this in the MACECalculator, we get the same results as before. Will merge this now. (Sorry, took a little longer for me to see this even though it is kind of obvious. What else besides the floating type precision should influence such a test when we do not change the MACE version ...)

[JaGeo]

@utf @janosh I have one question regarding backward compatibility: I decided to go with the new default (floating point precision of the model that one trained). This might influence the results a bit, but I think it is smarter than setting float32 by default. If we should change this, let me know.

[esoteric-ephemera]

@JaGeo good with adapting the test as needed but good to know it was a dtype issue