Add notebook that explains neural-lam pytorch tensors graph disk format

[leifdenby]

Describe your changes

Add notebook that describes the current neural-lam pickled pytorch tensors graph disk format. I feel like this documentation is currently missing, and since we write to this format in weather-model-graphs I think this would be a place where we could document it. This documentation could maybe also exist within neural-lam itself. But I feel like having this notebook could be a starting point at least.

NOTE: The notebook also shows that there is currently a bug in the index ordering in adjacency lists stored into pickle pytorch Tensors.

No change of dependencies needed.

Issue Link

n/a

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.

[leifdenby]

I've cleared the notebook output, but here's a rendered file from the most recent commit leifdenby@daa5f02:
saving_graphs_for_neural_lam.html

Creating these plots @joeloskarsson (see end of notebook) I think what is currently implemented when saving to neural-lam format actually has a bug, the one which you're fixing in https://github.com/mllam/weather-model-graphs/pull/34/files#diff-bdf567174bcb3d8da552f28644f860a91c5019a07b6238b7d85e35bd346c9edbL118. What do you think?

[leifdenby]

Adding link to comment @joeloskarsson made which I think is important here:

assumptions wrt node index order and subgraph order (in terms of m2g, g2m and m2m, and its levels) that neural-lam assumes?

• There is no assumption about the subgraph order, in the sense of m2g, g2m and m2m, as they are stored separately.

• In hierarchical mesh graphs the bottom level should come first, followed by higher levels in the hierarchy. But note that this is only the ordering of levels in the list, there are no requirements on one level having node indices higher/lower than any other.

• We are currently saying in neural-lam that mesh nodes must have the index before grid nodes (https://github.com/mllam/neural-lam/blob/4d8aceade489ff23f1a11cf0c9f9cd967902a35a/neural_lam/models/base_graph_model.py#L22), but this is not crucial and we should really change this as we integrate wmg better with neural-lam. I already did this on the research branch. See the point "Move all edge_index reindexing ..." in Enabling training on interior and boundary datastore neural-lam#138.

[joeloskarsson]

This should be very useful indeed! Looking at the output there seems to a bug yes. It could very well be the label sorting that's missing, that you linked to as fixed in #34. Have you tried rebasing this on top of that PR and see if you have the same issue?

[joeloskarsson]

grid-to-grid? What would that be? I think this is mentioned at other places in the notebook as well, but maybe just a typo.

[leifdenby]

yes, that was a typo 😳 thanks for catching this!

[leifdenby]

Have you tried rebasing this on top of that PR and see if you have the same issue?

No, I will try that and see :) thanks

[leifdenby]

Now that #34 is in main (with the sorting that you implemented @joeloskarsson) I merged main into my branch for this work.

The edges in the loaded files (neural-lam pickled pytorch tensor files) was previously:

Now it looks like this:

And here the edges from the networkx.DiGraph as defined by the node labels:

So I think there is still something going wrong since the mesh indexes for the m2m subgraph edges start at zero (of course this might not matter currently in neural-lam since we always offset the indexes so that the first index is 0, and then assume that the index values are consecutive - that's right, right?)

I think a bit more digging is needed here...

I've included the cell outputs in the notebook in the most recent commit (a7a8e9e).

[joeloskarsson]

Ok, I think that maybe we viewed different things as wrong. Or maybe I just did not pay close enough attention. I was mostly worried that not all mesh nodes seemed to be part of an edge in g2m. But maybe that was just how you set up the graph?

If I understand you correctly, you view the problem as m2m being saved with the first mesh node having index 0 in edge_index? But I don't think this is something you can get away from, if converting from networkx to pyg using pyg.utils. There are no grid nodes in the m2m graph, so naturally the first mesh node will have index 0. I suppose you could offset all mesh node indices with num_grid_nodes in the tensor, but that seems like a bit of a hack. But why is this a problem? As you wrote, neural-lam anyhow needs the m2m node indices to start with 0, and if they don't we now reindex them there.

[leifdenby]

There are no grid nodes in the m2m graph, so naturally the first mesh node will have index 0. I suppose you could offset all mesh node indices with num_grid_nodes in the tensor, but that seems like a bit of a hack. But why is this a problem? As you wrote, neural-lam anyhow needs the m2m node indices to start with 0, and if they don't we now reindex them there.

I had expected that the node-indices in the tensors saved to disk would be globally unique integers. Rereading what you've written here makes me realise we should document as part of the spec in mllam/neural-lam#323 the node-index values in the adjacency-list tensors and how they relate to the indices of the feature tensors (in effect what we are discussing here: mllam/neural-lam#323 (comment))

I think it would make most sense if the way we reference nodes in the edge_index values was if each nodeset (e.g. grid or mesh nodes) would start with the first node being at index value 0 and the last at N_grid-1 (for grid nodes) or N_mesh-1 (for mesh nodes). That isn't what is happening in the image I posted above (middle one) where I have loaded the tensor files produced by wmg. In that image for example the g2m edges the node-index for the mesh (g2m dst) nodes starts at 12 rather than 0.