Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ASE Calculator Updates #299

Merged
merged 8 commits into from
Oct 1, 2024
Merged

Conversation

melo-gonzo
Copy link
Collaborator

This pr updates the ASE calculator interface to accommodate a few changes:

  • Use models that were not trained with matsciml. This eliminates the requirement for the model to be one of ForceRegressionTask, etc.
  • Map from model output keys to keys that are expected by ASE. Currently, the calculator manually renames model outputs to expected outputs by ase, for example the model output force gets mapped into forces. This mapping should now be passes as a dictionary to the calculator on instantiation if required.
  • The calculator will now iterate through the properties passed into _calculate to build the results dictionary.
  • Add the concatenate_keys function which also adds properties to PyG graphs that are expected by models like MACE, add a few extra keys to _format_atoms, and change when type casting is called.
  • Updated tests to reflect new changes, added test using a pretrained matgl model.

matsciml/interfaces/ase/base.py Show resolved Hide resolved
# run the data structure through the model
output = self.task_module.predict(data_dict)
else:
output = self.task_module.predict(atoms)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So this requires other models to have a predict method as well?

If it's a model that isn't a matsciml model, I don't know if it's a better idea to just simply run their __call__ or forward method. If you don't think that's a good idea or have a reason for doing so, you should add in an assertion in the __init__ that checks that the input model has a predict function.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, actually I think just running their forward method would be good. Updated here 81f1ab0

matsciml/interfaces/ase/base.py Show resolved Hide resolved
matsciml/interfaces/ase/tests/test_ase_calc.py Outdated Show resolved Hide resolved
Copy link
Collaborator

@laserkelvin laserkelvin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, merge when ready

@melo-gonzo melo-gonzo merged commit dbba835 into IntelLabs:main Oct 1, 2024
3 checks passed
@melo-gonzo melo-gonzo deleted the calculator-refactor branch October 1, 2024 16:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants