Add a decoding mask option to only include subset of grid nodes in m2g

[joeloskarsson]

Describe your changes

For LAM forecasting it is reasonable to not always predict values in nodes that are considered part of the boundary forcing. This is in particular a change planned for neural-lam. When we consider the graphs involved, this means that the g2m edges should only connect to a subset of the grid nodes.

This PR introduces an option decode_mask that allows for specifying an Iterable of booleans (e.g. a numpy-array) specifying which of the grid nodes should be included in the decoding-part of the graph (m2g). This allows in the LAM case to specify such a mask with True for the inner region nodes.

This builds on #32, which should be merged first. Here is a diff for only this PR in the meantime: joeloskarsson/weather-model-graphs@general_coordinates...decoding_mask

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

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
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the documentation to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging)

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• author has added an entry to the changelog (and designated the change as added, changed or fixed)
• Once the PR is ready to be merged, squash commits and merge the PR.

[sadamov]

Looks good! Agree that this is required. Testing for "less than" seems sufficient 👍

More general comment: I am slightly worried about the mask being generated in mllam-data-prep and then being used here and in neural-lam. We probably need some solid checks in neural-lam about the dimensions of the graph, the data-tensors and the mask.

[joeloskarsson]

@leifdenby if you agree with my fixes this should now be done, except for the documentation notebook. The tests fail because of the existing test notebook docs/decoding_mask.ipynb. I did not want to remove that file in case you want to base the documentation on this, but it needs some cleanup and text to be a useful documentation.

[leifdenby]

Thank you for adding the tests and in general bringing this forward! I have added a few comments.

I have also updated the notebook locally on my computer. Do you want me to push to your branch? I haven't reworked the part of the notebook that describes the changes to the graph as stored to disk, but I would like to do that too. I think with your input my comments below I will be able to do that.

[joeloskarsson]

@leifdenby I fixed and responded to all your comments now.

Please commit the updated documentation notebook directly here and I can have a look.

[leifdenby]

Thanks for making all those changes! I will update the notebook now and then we can merge this in. Maybe you want to update the changelog now while I do that @joeloskarsson? Thanks 😊

[leifdenby]

I decided to complete remove the part of the notebook which examines what is actually written to disk @joeloskarsson. I did this so we can get this PR merged in, and improve the documentation later.

Wrt what I stored to disk, I instead I have started a notebook on the current neural-lam pickled pytorch tensors disk format here: #52. This has made me realise that I think there is a bug in the currently implementation (as I think you have also realised). If you have time I would really appreciate your thoughts on that PR too. Thank you

[joeloskarsson]

Thanks a lot for the notebook! I just made some small tweaks in 69f7f84.

I think this should be good to be merged now. Let me know if you have anything more or if I should go ahead.

[leifdenby]

Thanks a lot for the notebook! I just made some small tweaks in 69f7f84.

Love the tweaks ❤️

I think this should be good to be merged now. Let me know if you have anything more or if I should go ahead.

Yup, we just need a changelog entry and this will be ready to go in :)

[leifdenby]

Awesome! This means that I can release v0.3.0 since everything on that milestone is done https://github.com/mllam/weather-model-graphs/milestone/3 🥳 I will get onto that now.