megnet
megnet copied to clipboard
materialsvirtuallab
Can MEGNet be made end-to-end differentiable wrt atomic positions of atoms?
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 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.
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!
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.
@chc273 did you implement this in order to run the phonon calcs with m3gnet?
@CompRhys Do you mean energy derivative to calculate the forces? Yes. m3gnet will be a new codebase though
Yep, will await new codebase and see if it can be back ported. Is new code also TF based?