IntelLabs / matsciml

Open MatSci ML Toolkit is a framework for prototyping and scaling out deep learning models for materials discovery supporting widely used materials science datasets, and built on top of PyTorch Lightning, the Deep Graph Library, and PyTorch Geometric.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

[Feature request]: Refactor MACE to be compatible with task abstractions

laserkelvin opened this issue · comments

Feature/behavior summary

The current implementation of MACE is not fully compatible with the task abstraction we have designed: for example, force computation is hardcoded into the model similar to how the original OCP models were implemented. This means that it is not plug-and-play like the other models in being able to freely compose to an end-to-end pipeline.

Request attributes

  • Would this be a refactor of existing code?
  • Does this proposal require new package dependencies?
  • Would this change break backwards compatibility?
  • Does this proposal include a new model?
  • Does this proposal include a new dataset?
  • Does this proposal include a new task/workflow?

Related issues

N/A

Solution description

There are a number of things that need to be done in order to have MACE comply with the rest of the pipeline. To preserve backwards compatibility, it might be a better option to refactor the existing architecture as a "vanilla" implementation, and have a second compliant architecture. A "seamless" solution would be to have the MACE class act as a wrapper, which subsequently comprises both vanilla and conforming versions of the architecture.

A quasi-ordered task list (that will get updated) that should get converted into individual issues/PRs where applicable:

  • #116
  • Isolate/modularize unneeded utility functions, particularly those in support of independent training
  • Duplicate and rewrite MatSciML conforming MACE to use IrrepOutputBlock
  • #83
  • Register MACE in model registry

Additional notes

N/A

Closing this for now as it pertains to the ported version of MACE. With a wrapper based on upstream, we have a different set of things that need to be done, and will open a new issue tracking those.

Can reopen this if and when we are doing maintenance work for this version.