Add output_layer argument to M3GNet forward method to return intermediate layer outputs

[JiQi535]

Summary

In this PR, the forward method of M3GNet is modified so that intermediate outputs in "embedding", "gc_1", "gc_2", "gc_3", and "readout" layers can be extracted, while the "final" layer output stays as the default output. This allows the usage of M3GNet intermediate outputs as atom, bond, and structure features. Documentations and tests are provided.

Major changes:

1. Added the argument of output_layer into the forward method.
2. Added comprehensive tests for the predict_structure method.
3. Added the argument of output_layer into the predict_structure method.
4. Edited the variable names in the readout layer of M3GNet to be more precise. For example, the node_vec is renamed to field_vec, as this variable can either be node feature or edge feature instead of just node feature, depending on the argument field.

Todos

Follow-up analysis for the effectiveness of atom features.

Checklist

• [x] Google format doc strings added. Check with ruff.
• [x] Type annotations included. Check with mypy.
• [x] Tests added for new features/fixes.
• [ ] If applicable, new classes/functions/modules have duecredit @due.dcite decorators to reference relevant papers by DOI (example)

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Comparison is base (6759ffb) 98.59% compared to head (32b6f73) 98.61%.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #212 +/- ## ========================================== + Coverage 98.59% 98.61% +0.01% ========================================== Files 28 28 Lines 1927 1943 +16 ========================================== + Hits 1900 1916 +16 Misses 27 27 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[shyuep]

I think it is an extremely bad practice to put this in the predict_structure method. The predict_structure should return the expected target of the model. You should put it in a different method.

[kenko911]

Hi @JiQi535, thanks for your PR. I agree that we should have another class to extract features from structures instead of just putting in the function predict_structure, which is mainly for property prediction. Maybe we can call the new function "extract_features" and then we can also remove the unnecessary "final" option in the output_layer?

[JiQi535]

Thanks for the comments @shyuep @kenko911 . I have added a different method, featurize_structure, and the predict_structure method has been changed back to its original format to only return the expected target of the model.

By default, featurize_structure now returns the "readout" layer output and the "gc_3" layer output of intensive and extensive M3GNet models, respectively. Also, tests are updated, and a ValueError is raised if output_layer is not one of the allowed names in ["embedding", "gc_1", "gc_2", "gc_3", "readout", "final"].

[shyuep]

There is just way too much code and this is poor software design.

1. featurize_structure should simply return a dict of all intermediate and final layers. This way, there is no need to recompute every time you need one layer.
2. predict_structure should call featurize_structure in some way. Even better, the forward method actually stores the computed layers and you can access it immediately.

[JiQi535]

@shyuep The two suggestions have been realized in the updated code. My edits do not change the default output of forward method, so that other related codes, e.g., the Potential class, do not need change. This requires the addition of an argument in the forward method, i.e., return_all_layer_output. This implementation also does not increase the memory consumption of the forward method when return_all_layer_output is False as default.