feat(graphs): Introduce RadiusAreaMaskBuilder#1004
Conversation
for more information, see https://pre-commit.ci
|
Still need to do some cleanup (some documentations and typing need to be fixed), but feel free to have a look at it @JPXKQX |
for more information, see https://pre-commit.ci
|
Hi Michiel, this looks really good. As I see it, this could replace the |
| reference_node_name: str, | ||
| margin_radius_km: float = 100, | ||
| mask_attr_name: str | None = None, | ||
| use_gpu: bool = False, |
There was a problem hiding this comment.
Looking at the results, shouldn't we use cpu always? It seems to be faster in both cases: with and without torch-cluster
There was a problem hiding this comment.
I was thinking maybe later we can exploit the fact that we have multiple GPU available and devide the work over the different GPUs for even more speed up. Not sure if what I say make sense. I don't know if torch-cluster supports parallelising this kind of functions/work?
There was a problem hiding this comment.
As far as I know, only "single-device" parallelism is supported. We need something similar to the model sharding in anemoi-models (we should have the code to split into shards and gather the shards later). I think this will be easier to explore once we move the graph inside the model.
There was a problem hiding this comment.
Even in that case, I would expect this to happen automatically, right?
Not really, I would also be in favour of removing it. Just left it in for now if people would want to benchmark it themselves. |
|
Probably best to use this PR to add some additional tests for LAM or SG hidden meshes. |
2 participants
When building a Stretched Grid or LAM hidden mesh from a refined Icosahedron, the mesh structure is defined by using a
KNNAreaMaskBuilderfor calculating which hidden nodes fall under the high resolution part of the data mesh.This is currently done using a nearest neighbour search with
sklearn.neigbors.NearestNeighbors. When going to very high refinement levels (>10), this becomes a bottleneck in the graph building (~1h)This PR adds a new MaskBuilder-class:
RadiusAreaMaskBuilderthat uses a radius search either withtorch_geometric.nn.radiuswhentorch-clusteris installed or withscipy.spatial.cKDTreewhentorch-clusteris not available.This results in a 60% speedup. Some benchmarks:
numpy.float64DotProducttorch.float64DotProducttorch.float32torch.float64torch.float64numpy.float64DotProducttorch.float32torch.float64torch.float64numpy.float?Some figures confirming identical Stretched Grid Meshes (4,6) for the MEPS domain
As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/
By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.