Introduce graph classes for graph-based models

[joeloskarsson]

Describe your changes

The actual graph constructions (mesh graphs + m2g and g2m edges) currently take the form as a set of tensors in the code.
All these tensors are passed around and used in the different model classes in a somewhat unstructured way. See for example:

neural-lam/neural_lam/models/base_hi_graph_model.py

Lines 43 to 45 in 066efe0

	mesh_dim = self.mesh_static_features[0].shape[1]
	mesh_same_dim = self.m2m_features[0].shape[1]
	mesh_up_dim = self.mesh_up_features[0].shape[1]

This also means that there are some more or less hacky ways that we differentiate between hierarchical and "flat" graphs (see utils.load_graph):

neural-lam/neural_lam/utils.py

Line 111 in 066efe0

def load_graph(graph_name, device="cpu"):

The idea of this change is to keep the full representation of a graph together in one object. This means that all tensors (edge_index and features) related to the graph sits together in a coherent structure.

At the moment this loads the graph tensors from the same kind of saved .pt files as in the current implementation. There has been some discussion about just directly serializing these graph objects, and then using them also in https://github.com/mllam/weather-model-graphs.

Some TODOs:

• Try to turn into dataclasses, while also extending nn.Module. Test that moving to GPU still works.
• Implement hierarchical graph class
• Integrate graph classes to be used in the models

Issue Link

Does not directly solve an issue, but should be used in #4.

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)

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
• I have commented my code, particularly in hard-to-understand areas
• I have added tests that prove my fix is effective or that my feature works
• I have requested a reviewer and an assignee (assignee is responsible for merging)
• I have given the PR a name that clearly describes the change, written in imperative form (context.

Checklist for reviewers

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

• the code readable
• the code well tested
• the code documented
• the code easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change (in section
  reflecting type of change, for example "bug fixes", add section where
  missing)

Checklist for assignee

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