search

materialsvirtuallab / megnet

Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals
BSD 3-Clause "New" or "Revised" License
502 stars 156 forks source link

Can MEGNet be made end-to-end differentiable wrt atomic positions of atoms? #301

Open CompRhys opened 2 years ago

CompRhys commented 2 years ago

How hard would it be to change the construction of the local environment graph to be end-to-end differentiable? This would involve the neighbour list function inside pymatgen being refactored to include a sub-function that returns the index and image index allowing the distances to be computed from TF tensors of the atom positions?

Thinking about BOWSR there would be some benefit to regularising the force predictions of the MEGNet surrogate model to be ~0 given that it is trained on the relaxed structures in MP. As far as I understand the model is not currently given any signal about these structures being local minima.

CompRhys commented 2 years ago

link to #170

chc273 commented 2 years ago

@CompRhys For the pymatgen part, i have done the function you described. Try get_neighbor_list method of Structure.

End-to-end is definitely doable by starting from atom positions.

chc273 commented 2 years ago

For the idea on BOWS, what you described is correct. I think this is definitely worth trying by including the force regularization. Please keep me posted, thanks!

CompRhys commented 2 years ago

Am mostly a pytorch user but if the function is already implemented I will have a look at whether it's easy to make the change here after I submit my thesis in a few weeks. Conceivably the change should be hot swappable provided there are no issues arising out of differences in precision between numpy arrays and TF tensors.

CompRhys commented 2 years ago

@chc273 did you implement this in order to run the phonon calcs with m3gnet?

chc273 commented 2 years ago

@CompRhys Do you mean energy derivative to calculate the forces? Yes. m3gnet will be a new codebase though

CompRhys commented 2 years ago

Yep, will await new codebase and see if it can be back ported. Is new code also TF based?

chc273 commented 2 years ago

@CompRhys thanks. Yes, the code is in TF.