Separate interior state and boundary forcing to only predict state

[joeloskarsson]

Describe your changes

The goal of this PR is to establish a clear separation of the (predicted) state in the interior region and the boundary forcing coming from outside (and potentially overlapping with) the limited area. This PR performs these changes on the modeling side. That is from each batch is fetched from the dataset class and onward. Including how the tensors are propagated through the model, loss calculation, evaluation and plotting. This should be complemented with a separate PR for handling the data-loading side of things, where the boundary forcing could come from a separate dataset. That change should then build upon #66, once merged.

Note: In order to allow for working on this before the change has been done on the data loading side this currently includes changes in the MEPS npy dataset class that separates state and boundary already in there. This defines the interface between the dataset and model (currently missing #64 from #66, but that can easily be added later) and allows for working on these separately.

After this change:

• The model will only predict outputs within the limited area considered
• Plots will only include points within the limited area (this could be expanded in a future PR to plot also boundary fields, but would require a mapping between state and boundary forcing dimensions to plot together)
• Graphs will have to be created using a boundary mask as in Add a decoding mask option to only include subset of grid nodes in m2g weather-model-graphs#34 to make sure that the g2m-component only maps to the interior nodes.

Dependencies:

This introduces a dependency to https://github.com/mllam/weather-model-graphs. In particular, this dependency should be adjusted before merging to require a version after mllam/weather-model-graphs#34 has been merged.

Issue Link

No issue specific to the separation of interior state. This relates to the overall rework of Reading Training Data, but would be good to put as separate point on roadmap.

This includes graph-creation with wmg (#83).

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 README 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). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

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.

[joeloskarsson]

This is now ready for a first review. As mentioned in the description this is only changes to the modeling side of things, and I have tweaked the existing MEPS dataloading to be able to test these changes. We will have to see if the changes to dataloading (building on top of #66) should be in a separate PR or just added to this one. I could see both as good solutions.

There are two things preventing tests to pass for this:

• This changes the grid_shape_state entry in the config to only mean the interior, not boundary. Thus the grid_shape_state in the reduced MEPS data used for tests is wrong. How/if we want to change this depends on how we want to go about with the ordering of this and Add "datastores" to represent input data from zarr, npy, etc #66.
• This introduces graph creation using weather-models-graph, and in particular relies on the masking functionality in Add a decoding mask option to only include subset of grid nodes in m2g weather-model-graphs#34. This needs to be merged for the tests to construct the correct graph.

My idea is that this could get a quick review right now, just considering the current changes, and then we can make up a plan w.r.t. the merging or continued work on this. Don't spend time on the changes to the MEPS dataloading made here, as that will anyhow be replaced with #66.

[joeloskarsson]

Making a new PR for this, as this is becoming a bit bloated and this is now rebased on #90 .

[leifdenby]

For future reference this is replaced with #93