SBRA-739 redefine the graph creation method

[A669015]

This PR refactor the code to avoid the computation of the graph topology for each sample.
The graph topology is now computed on the datamodule side at initialization, and passed to the ligthning module through the argument linking feature of ligthning. Then at training time, the dataloader only send the features to the model.

[github-actions]

Coverage report for weather-forecast/gravity-wave-drag/cnns model project

File	Coverage		Missing
All files	100%	✅

Minimum allowed coverage is 0%

Generated by 🐒 cobertura-action against db03d17

[github-actions]

Coverage report for reactive-flows/cnf-combustion/unets model project

File	Coverage		Missing
All files	97%	✅
data.py	97%	✅	134 148
unet.py	96%	✅	95 125

Minimum allowed coverage is 0%

Generated by 🐒 cobertura-action against db03d17

[github-actions]

Coverage report for reactive-flows/cnf-combustion/gnns model project

File	Coverage		Missing
All files	94%	✅
data.py	93%	✅	82 160 214 222 239 245 317-321
inferer.py	89%	✅	214-240
plotters.py	94%	✅	156-157 160-161 182 191

Minimum allowed coverage is 0%

Generated by 🐒 cobertura-action against db03d17

[github-actions]

Coverage report for weather-forecast/ecrad-3d-correction/unets model project

File	Coverage		Missing
All files	96%	✅
data.py	98%	✅	97
dataproc.py	93%	✅	164 168-170 215-216
models.py	98%	✅	52
unet.py	94%	✅	92 122 125

Minimum allowed coverage is 0%

Generated by 🐒 cobertura-action against db03d17

[mikael10j]

Just a quick review

[cbovalo]

Just a few comments:

• in the pyg_data object, I suggest setting edge_index to None to make it clearer that the adjacency matrix will be given later by graph_topology,
• I think the egde_index should be a buffer of the nn.Module to be put in the correct device automatically (right now it stays on cpu),
• I try to think to a better way to pass the grid shape (for plotting) instead of using the position. I mean do we need a tensor of shape [num_nodes, num_dimensions] to get only 3 numbers ?

[A669015]

• I think the egde_index should be a buffer of the nn.Module to be put in the correct device automatically (right now it stays on cpu),

Do you suggest to register the edge_index tensor in the state_dict, so in the model checkpoint ? We can make it non-persistent to keep the checkpoint as light as possible.

• I try to think to a better way to pass the grid shape (for plotting) instead of using the position. I mean do we need a tensor of shape [num_nodes, num_dimensions] to get only 3 numbers ?

What do you mean by "a tensor of shape [num_nodes, num_dimensions] to get only 3 numbers" ? The grid_shape is a tuple of 3 values.

[cbovalo]

• I think the egde_index should be a buffer of the nn.Module to be put in the correct device automatically (right now it stays on cpu),

Do you suggest to register the edge_index tensor in the state_dict, so in the model checkpoint ? We can make it non-persistent to keep the checkpoint as light as possible.

Not in the state_dict (setting persistent=False in register_buffer)

• I try to think to a better way to pass the grid shape (for plotting) instead of using the position. I mean do we need a tensor of shape [num_nodes, num_dimensions] to get only 3 numbers ?

What do you mean by "a tensor of shape [num_nodes, num_dimensions] to get only 3 numbers" ? The grid_shape is a tuple of 3 values.

Exactly! What I mean is that you add to graph_topology an attribute pos which is a tensor with shape (num_nodes, 3) - Cartesian coordinates in 3D. You retrieve the min/max in each axis to get the domain size. Maybe this shape could be returned by the Datamodule and be liked to the LightningModule.