Add (Structure|Molecule).calculate() + Molecule.relax(), improve Structure.relax()

[Andrew-S-Rosen]

Closes #3042. Closes #3043. Closes #3047.

Links:

• Wrong typehint in Structure.relax #3042
• Add a Structure.calculate() feature analagous to Structure.relax #3043
• Structure.relax only works if m3gnet is installed even if it is not used #3047

In addition to the prior issues, I also increased the range of optimizers that can be passed to Structure.relax so the user doesn't have to rely solely on FIRE.

For molecules, I added a Molecule.relax() feature that defaults to GFN2-xTB and an analogous Molecule.calculate().

[shyuep]

I suggest:

1. We have a run_calculation method that replaces and extends most of the code in relax to be more flexible.
2. We have relax call run_calculator with a specific set of parameters to do a relaxation.

Structure.run as a naming sounds a bit strange to me.

[Andrew-S-Rosen]

Agreed. That'll clean things up and prevent duplication. And that's a much better name. That way it doesn't sound like the structure is running away...

[Andrew-S-Rosen]

@shyuep, @janosh: This is ready for either of you to review.

A couple of notes:

• I did not see any clean way of using Structure.run_calculation as the base logic for Structure.relax. They both do fairly different things, and the code between the two reflects that. Nonetheless, I did use .run_calculation instead of run as requested.
• I can add in the analogous functions in Molecule once things look good for Structure. I imagine we should not use m3gnet as the default for Molecule, so I'm open to suggestions. This might also be a good spot to consider refactoring since we don't necessarily want to copy/paste things between Structure and Molecule. They'll be basically the same except for no options to relax the cell in the latter.
• For some reason, m3gnet gives slightly different lattice constants on Mac/Windows than on Linux for one of the tests. I can modify the assertion, but I want to alert you of this odd behavior.
• Instead of using the m3gnet custom TrajectoryObserver, I have switched to using ASE's native Trajectory object when return_trajectory is True in Structure.relax. The latter has all the desired properties (e.g. energies, forces, cells, stresses, etc.) without requiring the m3gnet dependency for all ASE calculations.
• I will need a bit of help with the mypy test failures.

[janosh]

That was fast! Thanks @arosen93. Sorry for not bringing this up earlier but what do you think of Structure.calc() instead of Structure.run_calculation()?

For some reason, m3gnet gives slightly different lattice constants on Mac/Windows than on Linux for one of the tests. I can modify the assertion, but I want to alert you of this odd behavior.

That might be a TensorFlow quirk. I hope/expect the new matGL M3GNet built using DGL doesn't have that. So maybe not worth it to dig too deeply for now.

[janosh]

@arosen93 Sorry for the slow reply! This is looking great! I left 2 questions/comments. Also, I saw that _prep_calculator and _relax are marked as private but since they have user-facing features, i.e. raise error messages, I think it would be good to add a test for that behavior. Can try to do that later.

[Andrew-S-Rosen]

@janosh --- ready for either re-review or merging. Some of the Mac tests died but seems to be a problem for everyone.

[janosh]

Thanks @arosen93! Took a few iterations but this was no small PR and I think it turned out really well! 😎 Thanks for the excellent test coverage! 🧪