Stop hardcoding names of Cartesian coords

[leifdenby]

Instead of assuming that the regular grid Cartesian coordinates are called x and y use the coordinate names from the datastore config of the coordinates that were stacked.

Describe your changes

This PR removes the hardcoding of x and y as the names of the Cartesian coordinates in MDPDatastore. This means that if the regular grid coordinates are not called x and y but instead for example the training data is on a regular lat/lon grid (i.e. 2D coordinate dimensions are for example called lat and lon), neural-lam will still be able to use this training data.

This came about due to the discussion on mllam/mllam-data-prep#45 (comment) where @joeloskarsson was trying to use ERA5 data where the dimension of the variables are (time, latitude, longitude, ...) i.e. the regular grid 2D spatial coordinates are not called x and y, instead latitude and longitude. This fix will allow for MDPDatastore in neural-lam to work with this ERA5 dataset, and any other dataset with differently named coordinates.

I am realising that calling these coordinates "Cartesian" is not strictly correct. Should we call them "regular grid coordinates", or what is the better name here?

No change of dependencies needed for this fix.

Issue Link

mllam/mllam-data-prep#45 (comment)

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]

Maybe we can call it grid_coordinates or spatial_coordinates? I also don't get why CARTESIAN_COORDS was in all caps, as this is not a constant.

Other than the variable name this looks great :)

[joeloskarsson]

(I realize I never added a review properly here)

Let's just change the variable names as discussed above, then this is good :)

[joeloskarsson]

Hmm, something seems to go wrong with the dummy datastore in the tests now

[leifdenby]

Hmm, something seems to go wrong with the dummy datastore in the tests now

I hadn't changed the variable defining the spatial coordinates within the dummy datastore 😊
Also, I realised the implementation we currently have on main wrt to unstacking might cause the where the spatial coordinates appear in the dimension order might change. The new implementation handles that to ensure the spatial coordinate stay in the same order as where the grid_index was (before unstacking)

[leifdenby]

The tests are passing now so this can go in v0.6.0 I think if you can check again @joeloskarsson and approve if you're happy with it :)

[joeloskarsson]

Looks good to me now! Please remember to add a changelog entry before merging.