neural_lam.utils.graph#

Loading and zero-indexing of mesh/grid graph tensors.

Attributes#

LEGACY_GRAPH_SPEC_VERSION

Functions#

`compute_grid_input_dim`(→ int)	Compute the total grid input dimensionality of a graph step predictor.
`load_and_register_graph`(→ bool)	Load a graph and register its tensors on `module`.
`load_graph`(→ tuple[bool, dict[str, Any]])	Load all tensors representing the graph from graph_dir_path.
`zero_index_edge_index`(→ torch.Tensor)	Make both sender and receiver indices of edge_index start at 0.
`zero_index_g2m`(→ torch.Tensor)	Zero-index the g2m (grid-to-mesh) edge index, or undo this operation.
`zero_index_m2g`(→ torch.Tensor)	Zero-index the m2g (mesh-to-grid) edge index, or undo this operation.

Module Contents#

neural_lam.utils.graph.compute_grid_input_dim(datastore: neural_lam.datastore.BaseDatastore, num_past_forcing_steps: int, num_future_forcing_steps: int) → int[source]#

Compute the total grid input dimensionality of a graph step predictor.

The grid input concatenates the two previous states, the grid static features and the windowed forcing (past + current + future forcing steps).

Parameters:

datastore (BaseDatastore) – Datastore providing the number of state, static and forcing variables.
num_past_forcing_steps (int) – Number of past forcing steps included in the input window.
num_future_forcing_steps (int) – Number of future forcing steps included in the input window.

Returns:

Total grid input dimensionality.

Return type:

int

neural_lam.utils.graph.load_and_register_graph(module: torch.nn.Module, datastore: neural_lam.datastore.BaseDatastore, graph_name: str, mesh_node_features_scaling: float) → bool[source]#

Load a graph and register its tensors on module.

Loads the graph graph_name from the datastore’s graph directory via load_graph(), then registers each tensor as a non-persistent buffer and each non-tensor (e.g. BufferList) as a plain attribute on module.

Parameters:

module (torch.nn.Module) – Module to register the graph tensors and attributes on, in place.
datastore (BaseDatastore) – Datastore whose root_path holds the graph directory.
graph_name (str) – Name of the graph directory (under <root>/graph) to load.
mesh_node_features_scaling (float) – Scalar used to normalize mesh node coordinate features for graphs in the current on-disk format; forwarded to load_graph().

Returns:

Whether the loaded graph is hierarchical.

Return type:

bool

neural_lam.utils.graph.load_graph(graph_dir_path: str | pathlib.Path, mesh_node_features_scaling: float, device: str = 'cpu') → tuple[bool, dict[str, Any]][source]#

Load all tensors representing the graph from graph_dir_path.

Needs the following files for all graphs: - m2m_edge_index.pt - g2m_edge_index.pt - m2g_edge_index.pt - m2m_features.pt - g2m_features.pt - m2g_features.pt - mesh_features.pt

And in addition for hierarchical graphs: - mesh_up_edge_index.pt - mesh_down_edge_index.pt - mesh_up_features.pt - mesh_down_features.pt

Parameters:

graph_dir_path (str) – Path to directory containing the graph files.
mesh_node_features_scaling (float) – Scalar used to normalize mesh node coordinate features for graphs in the current on-disk format.
device (str) – Device to load tensors to.

Returns:

hierarchical (bool) – Whether the graph is hierarchical.
graph (dict) – Dictionary containing the graph tensors, with keys as follows: - g2m_edge_index - m2g_edge_index - m2m_edge_index - mesh_up_edge_index - mesh_down_edge_index - g2m_features - m2g_features - m2m_features - mesh_up_features - mesh_down_features - mesh_static_features

neural_lam.utils.graph.zero_index_edge_index(edge_index: torch.Tensor) → torch.Tensor[source]#

Make both sender and receiver indices of edge_index start at 0.

Parameters:: edge_index (torch.Tensor) – Edge index tensor of shape (2, num_edges).
Returns:: Edge index tensor with indices starting at 0.
Return type:: torch.Tensor

neural_lam.utils.graph.zero_index_g2m(g2m_edge_index: torch.Tensor, mesh_static_features: list[torch.Tensor], mesh_first: bool, restore: bool = False) → torch.Tensor[source]#

Zero-index the g2m (grid-to-mesh) edge index, or undo this operation.

Special handling is needed since not all mesh nodes may be present.

Parameters:

g2m_edge_index (torch.Tensor) – Edge index tensor of shape (2, num_edges).
mesh_static_features (list of torch.Tensor) – Mesh node feature tensors.
mesh_first (bool) – If True, mesh nodes are indexed before grid nodes.
restore (bool) – If True, undo zero-indexing (restore original indices).

Returns:

Edge index tensor with zero-based or restored indices.

Return type:

torch.Tensor

neural_lam.utils.graph.zero_index_m2g(m2g_edge_index: torch.Tensor, mesh_static_features: list[torch.Tensor], mesh_first: bool, restore: bool = False) → torch.Tensor[source]#

Zero-index the m2g (mesh-to-grid) edge index, or undo this operation.

Special handling is needed since not all mesh nodes may be present.

Parameters:

m2g_edge_index (torch.Tensor) – Edge index tensor of shape (2, num_edges).
mesh_static_features (list of torch.Tensor) – Mesh node feature tensors.
mesh_first (bool) – If True, mesh nodes are indexed before grid nodes.
restore (bool) – If True, undo zero-indexing (restore original indices).

Returns:

Edge index tensor with zero-based or restored indices.

Return type:

torch.Tensor

neural_lam.utils.graph.LEGACY_GRAPH_SPEC_VERSION = 'legacy'#