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: https://github.com/joeloskarsson/weather-model-graphs/compare/general_coordinates...decoding_mask

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)
• [ ] 📖 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).
• [x] I have performed a self-review of my code
• [x] For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• [x] 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
• [x] I have added tests that prove my fix is effective or that my feature works
• [x] I have given the PR a name that clearly describes the change, written in imperative form (context).
• [x] 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:

• [x] the code is readable
• [x] the code is well tested
• [x] the code is documented (including return types and parameters)
• [x] 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.

[joeloskarsson]

I agree that there is a need for checks w.r.t. the dimensions of the graph and feature tensors. The work I started on https://github.com/mllam/neural-lam/tree/graph_classes does contain quite a lot of this that I think will be useful, e.g.

• https://github.com/mllam/neural-lam/blob/c097c582a3329c405ffc0879e7c27582d6735ab1/neural_lam/graphs/base_weather_graph.py#L101-L118
• https://github.com/mllam/neural-lam/blob/c097c582a3329c405ffc0879e7c27582d6735ab1/neural_lam/graphs/flat_weather_graph.py#L35-L49.

[joeloskarsson]

This PR is now mainly waiting on #32 to be merged

[joeloskarsson]

To get this to actually work with neural-lam (or in particular with https://github.com/mllam/neural-lam/pull/84) I had to introduce some sorting of nodes and sub-graphs after splitting. Otherwise I was getting arbitrary node indices.

@sadamov @leifdenby could you please have a look at https://github.com/mllam/weather-model-graphs/pull/34/commits/350e0c0b8166594f57d20aabc39c92d59d72b82b and check that you don't see any problem with this.

To clarify here: We don't need all nodes to be in some very specific ordering. What we do need is that nodes in one specific sub-graph (e.g. "mesh-level 1") all have subsequent indices.

[sadamov]

@joeloskarsson to me this restriction seems justified. It's also an easy to understand requirement for folks who want to create their graphs outside of WMG. Just make sure to add it to both README's here and in NL. (and thanks for your work on this! :handshake:)

[joeloskarsson]

Good suggestions all around @SimonKamuk ! Was there something specific introduced by this PR that you think should be added to the readme? I was thinking that the example in the readme can continue to use default parameters, which is decode_mask=None.

Otherwise this is the plan for this PR now:

• [x] Make sure v0.2.0 is released
• [ ] Make sure #32 is merged
• [ ] Add this change to changelog, under unreleased to be in v0.3.0
• [ ] Merge

[SimonKamuk]

Well, it's not strictly necessary, and I think it's actually better suited for the documentation notebooks rather than the readme, but I just thought it would be nice to describe this feature with an example to make it easier for people to find. But again, it might not be necessary 😄

[joeloskarsson]

I do agree though that it could be nice to have a documentation notebook with an example of using this. Will see if I can put together something small.