[Feature Request]: interface with torch geometric

[bkmi]

Email (Optional)

No response

Problem

Machine learning researchers often keep data in the pytorch geometric format. Applying CHGNet to this data is currently not easily supported.

Furthermore, batch evaluation of the network is extremely unwieldy as it is, this would help address that issue.

Proposed Solution

I would propose to add a method to CHGNet which would work something like CHGNet.predict_geometric(data: torch_geometric.data.Data) -> CHGNetOutput.

Alternatives

Another option would be to introduce a converter from pytorch geometric format to your format.

Code of Conduct

• [X] I agree to follow this project's Code of Conduct

[BowenD-UCB]

Can you please elaborate the background of this? Do you have pyG graphs that you want to predict with CHGNet? Is there a reason why the original material (i.e. ase atoms or pymatgen structure) can not be used?

[bkmi]

I'm interested in computing gradients on PyG graphs w.r.t. CHGNet predictions. This would be pretty similar to a relaxation, but then I could do it in batches. (necessary for my purposes)

[janosh]

(necessary for my purposes)

Please elaborate.