The idea of this change is to keep the full representation of a graph together in one object. This means that all tensors (edge_index and features) related to the graph sits together in a coherent structure.
At the moment this loads the graph tensors from the same kind of saved .pt files as in the current implementation. There has been some discussion about just directly serializing these graph objects, and then using them also in https://github.com/mllam/weather-model-graphs.
Some TODOs:
[ ] Try to turn into dataclasses, while also extending nn.Module. Test that moving to GPU still works.
[ ] Implement hierarchical graph class
[ ] Integrate graph classes to be used in the models
Issue Link
Does not directly solve an issue, but should be used in #4.
Type of change
[ ] 🐛 Bug fix (non-breaking change that fixes an issue)
[x] ✨ New feature (non-breaking change that adds functionality)
[ ] 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
Checklist before requesting a review
[ ] My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
[ ] I have performed a self-review of my code
[ ] I have commented my code, particularly in hard-to-understand areas
[ ] I have added tests that prove my fix is effective or that my feature works
[ ] I have requested a reviewer and an assignee (assignee is responsible for merging)
[ ] I have given the PR a name that clearly describes the change, written in imperative form (context.
Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
[ ] the code readable
[ ] the code well tested
[ ] the code documented
[ ] the code easy to maintain
Author checklist after completed review
[ ] I have added a line to the CHANGELOG describing this change (in section
reflecting type of change, for example "bug fixes", add section where
missing)
Checklist for assignee
[ ] PR is up to date with the base branch
[ ] the tests passing
[ ] author has added an entry to the changelog (and designated the change as added, changed or fixed)
Once the PR ready to be merged, squash commits and merge the PR.
Describe your changes
The actual graph constructions (mesh graphs + m2g and g2m edges) currently take the form as a set of tensors in the code. All these tensors are passed around and used in the different model classes in a somewhat unstructured way. See for example: https://github.com/mllam/neural-lam/blob/066efe0a77f1ed741d9eb2f303ee920927d79ad2/neural_lam/models/base_hi_graph_model.py#L43-L45
This also means that there are some more or less hacky ways that we differentiate between hierarchical and "flat" graphs (see
utils.load_graph): https://github.com/mllam/neural-lam/blob/066efe0a77f1ed741d9eb2f303ee920927d79ad2/neural_lam/utils.py#L111The idea of this change is to keep the full representation of a graph together in one object. This means that all tensors (edge_index and features) related to the graph sits together in a coherent structure.
At the moment this loads the graph tensors from the same kind of saved
.ptfiles as in the current implementation. There has been some discussion about just directly serializing these graph objects, and then using them also in https://github.com/mllam/weather-model-graphs.Some TODOs:
nn.Module. Test that moving to GPU still works.Issue Link
Does not directly solve an issue, but should be used in #4.
Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
Checklist for assignee