Support anemoi.graphs #2
Conversation
|
|
There was a problem hiding this comment.
Thanks for the PR.
I have left some comments below, regarding naming of variables, arranging of code, and to comment and explain some parts of the code.
Additionally, I think, we should be careful with the use of get_attr and set_attr, because they have previously prevented serialisation of the model (and are generally really hard to decipher.)
| self._graph_name_hidden = config.graph.hidden | ||
| self._graph_name_hidden = config.graphs.hidden_mesh.name | ||
| self._graph_mesh_names = [name for name in graph_data if isinstance(name, str)] | ||
| self._graph_input_meshes = [ |
There was a problem hiding this comment.
These need a comment to explain what is being done here.
Probably should be a method too that takes graph_data as an input and processes it.
There was a problem hiding this comment.
From reading the code all the way to the end, it seems that these are filtering the names / keys of the graphs. This should be reflected in the variable names.
| sub_graph=self._graph_data[(self._graph_name_hidden, "to", self._graph_name_hidden)], | ||
| src_grid_size=self._hidden_grid_size, | ||
| dst_grid_size=self._hidden_grid_size, | ||
| sub_graph=self._graph_data.get((self._graph_name_hidden, "to", self._graph_name_hidden), None), |
There was a problem hiding this comment.
What happens if None is passed to the sub_graph?
There was a problem hiding this comment.
It will fail unless the TransformerProcessor is specified.
| self.trainable_hidden_size = config.model.trainable_parameters.hidden | ||
|
|
||
| def _register_latlon(self, name: str, key: str) -> None: | ||
| self.num_nodes = {name: self._graph_data[name]["coords"].shape[0] for name in self._graph_mesh_names} |
There was a problem hiding this comment.
The use of self._graph_mesh_names is redundant here, should probably be something like {name: graph["coords"].shape[0] for name, graph in self._graph_data.items()}.
There was a problem hiding this comment.
self._graph_data has both nodes (they key is a string) and connections (the key is a tuple of strings) in different items. It would be something like {name: graph["coords"].shape[0] for name, graph in self._graph_data.items() if isinstance(name, str)}. Do you think is it still clearer than with self._graph_mesh_names?
| # Register lat/lon | ||
| self._register_latlon("data", self._graph_name_data) | ||
| self._register_latlon("hidden", self._graph_name_hidden) | ||
| for name in self._graph_mesh_names: |
There was a problem hiding this comment.
This could just be in self._graph_data for simplicity. name should probably be renamed into something like mesh_key or grid_key.
There was a problem hiding this comment.
and mesh_name? do you think is it a good option?
| self._graph_name_data = config.graph.data | ||
| self._graph_name_hidden = config.graph.hidden | ||
| self._graph_name_hidden = config.graphs.hidden_mesh.name | ||
| self._graph_mesh_names = [name for name in graph_data if isinstance(name, str)] |
There was a problem hiding this comment.
What does this filter for?
JPXKQX
force-pushed
the
feature/graph_refactor
branch
from
305c940
to
dc0dc7d
Compare
Description
Support for new graph mappings:
Type of change
Please delete options that are not relevant.
Checklist before requesting a review