feat: new interactive graph visualization

[frazane]

Description

Adds a new lightweight way to visualize graphs interactively. No added dependencies (the javascript code runs entirely in the broswer and sources libraries via CDN). The way this is implemented is quite simple: the data is embedded directly in the HTML in json format using a jinja template. Then javascript code does the rest (yes, definitely with the help of coding assistants).

The largest graph on which this was tested was a cutout graph with n320 grid globally and a 2km grid over central Europe (the one we currently have at MeteoSwiss). Loaded in about 6s and then still runs relatively smoothly. Still, we might be reaching a limit with >1'500'000 nodes with lots of connections - we might need to find some optimizations to the code at some point as I am sure there are some left.

Nothing changes in terms of interface - this will simply add another output to the command:

anemoi-graphs inspect graph.pt output_dir/

Features

• fast GPU-based rendering
• full 3D scene navigation
• interactive controls for nodes and edges (toggles, radius scaling, colors, etc.)
• nodes selection and highlighting
• globe map
• only for nodes coordinates and edges (no attributes values, for now)

Graphs tested

• n320 global (+ hidden)
• CERRA LAM graph from RMI (+ hidden)
• COSMO SGM graph from MeteoSwiss (+hidden)
• n320 global hierarchical (+ multiple hidden levels)

---

---

📚 Documentation preview 📚: https://anemoi-training--609.org.readthedocs.build/en/609/

---

📚 Documentation preview 📚: https://anemoi-graphs--609.org.readthedocs.build/en/609/

---

📚 Documentation preview 📚: https://anemoi-models--609.org.readthedocs.build/en/609/

[mchantry]

Looks really cool! Thanks for the contribution @frazane

[anaprietonem]

@frazane this is really cool! Are you still planning to add more features? I'd be happy to review or test this if it helps!

[frazane]

@frazane this is really cool! Are you still planning to add more features? I'd be happy to review or test this if it helps!

Thanks! I think for this PR I would only make sure that this works for arbitrary graphs, including stretched grid graphs with many nodes (let's see how well this scales) and hierarchical graphs.

In subsequent PRs I see a few additional features we might add:

• extend selection logic to highlight not only the first next neighbours, but to arbitrary levels of connectivity (could be useful for hierarchical graphs to visualise the perception field)
• visualise arbitrary nodes and edges attributes
• (maybe, I am a dreamer) visualise model features on latent meshes extracted from a forward pass
• toggle to highlight orphan nodes

It would be great if you could test it (you can simply run the module with the temporary CLI and provide the path to a graph). Generally, the more graphs we throw at it, the better. Feedback is also very welcome! :)

[dietervdb-meteo]

Hi @frazane , is this always plotting things on a globe? LAM graphs for example cover only a limited area, would they get rendered (on part of the globe?) in the current implementation? (Don't have time to test now, can send you a LAM graph in case you would have time, in any case this comment should not stop you merging/developping this PR)

[frazane]

Hi @frazane , is this always plotting things on a globe? LAM graphs for example cover only a limited area, would they get rendered (on part of the globe?) in the current implementation? (Don't have time to test now, can send you a LAM graph in case you would have time, in any case this comment should not stop you merging/developping this PR)

Yes, it should work on arbitrary graphs. I tried with the stretched grid graph and it also worked. Can you send me a small LAM graph that I can test?

[fprill]

Very nice feature!

Just a quick note: I had to adapt the script slightly on my system.

@@ -82,14 +82,14 @@ if __name__ == "__main__":
     nodes, edges = load_graph(args.path, nodes=args.nodes, edges=args.edges)
 
     for node_set in nodes:
-        node_lats, node_lons = coords_to_latlon(nodes[node_set].numpy())
+        node_lats, node_lons = coords_to_latlon(nodes[node_set].to('cpu').numpy())
         nodes[node_set] = to_nodes_json(node_lats, node_lons, prefix=node_set)
 
     for edge_set in edges:
         src_nodes, dst_nodes = edge_set.split("_to_")
         src_names = [f"{src_nodes}_{i}" for i in range(len(nodes[src_nodes]))]
         dst_names = [f"{dst_nodes}_{i}" for i in range(len(nodes[dst_nodes]))]
-        edges[edge_set] = to_edges_json(src_names, dst_names, edges[edge_set].numpy().T)
+        edges[edge_set] = to_edges_json(src_names, dst_names, edges[edge_set].to('cpu').numpy().T)
 

[frazane]

Hi @frazane , is this always plotting things on a globe? LAM graphs for example cover only a limited area, would they get rendered (on part of the globe?) in the current implementation? (Don't have time to test now, can send you a LAM graph in case you would have time, in any case this comment should not stop you merging/developping this PR)

Yes, it should work on arbitrary graphs. I tried with the stretched grid graph and it also worked. Can you send me a small LAM graph that I can test?

Tested and works correctly 👍

[anaprietonem]

It would be great if you could test it (you can simply run the module with the temporary CLI and provide the path to a graph). Generally, the more graphs we throw at it, the better. Feedback is also very welcome! :)

Hey Francesco, I run for global gt and transformer graphs and it worked very nicely!

Those are o96, I will check also n320. But it was very easy to use and it loads very fast! so big kudos!. Some comments (mostly minor but in case you'd like some feedback):

• In the case of the transformer, if the edges are not passed and the default ones are used AttributeError: 'EdgeStorage' object has no attribute 'edge_index' the error AttributeError: 'EdgeStorage' object has no attribute 'edge_index' is not very easy to understand so we could have a check and raise a mistmatch before?
• if one wanted to change the colors of the edges to distinguish between data-hidden, and hidden-data, I guess for now they could just hardcode the list of colors? but maybe this could also be optional for users to overwrite

[jlcasador]

It is a wonderful tool. We have tested it using our stretched-grid model over the Canary Islands (2.5 km resolution). Some comments, just in case they are useful:

• It works and loads faster than the old tool, though the HTML files are still large when the resolution is very high, it needs time to load and it sometimes freezes. We also had this problem with the old tool, but it had the option --area to select a specific area to display, and that allowed us to lighten the HTML file. Would it be posible to implement this option too in a future version?
• The 3D perspective sometimes makes difficult to see the connections between nodes. I guess this is unavoidable, and the 2D tool will continue to be necessary in some cases.
• It seems that displaying 2.5 km resolution data is close to the limits of the tool. We have managed to zoom a lot (see the figure), though in that case the background map is a bit blurry (I guess a more detailed map might be used in that case).

[frazane]

Thanks @anaprietonem and @jlcasador for the feedback!

ana:

• Thanks, I did not try it with the transformer. I wonder if we could have then already some defaults for each model type? In any case, for now I'll make sure to raise a more informative error.
• Yes, this would be a nice addition.

jlcasador:

• Yes, definitely selecting a subset of the data points is a good idea - I will make sure to implement it. As you pointed out with large graphs the size of the HTML is pretty large, because the data itself is embedded in the file. So subsetting will help. As for the lagging, I suppose it depends on your hardware...I tried with a 2km mesh and it worked without freezes.
• Not sure what you mean here. If you select a node, can you not clearly see all the edges and nodes connected to it? But yes in any case, for sure some visualizations will be easier to understand in 2D.
• Yes, I would say 1-2km is currently getting close to the limits of the tool (if one wants to plot everything, but with the subsetting discussed above this will become less of an issue). The good thing is that there are likely still optimizations left to apply - I am not an expert in THREE.js at all. Would be nice to find someone who is.

[JPXKQX]

Hi Francesco, sorry for the delayed review. I just have one minor comment about the location: could we put it in interactive_3d_html.py and rename the previous one to interactive_2d_html.py, and use the current CLI from inspect.py? This would allow us to reuse some GraphInspector functionality, such as loading the graph or subsetting it with the area argument.

Also, I would set the default edges to 'None' and iterate over 'graph.edge_types', which will return all the edges already defined in the graphs.

Thanks for the contribution, really cool! Looking forward to have it merged

[frazane]

@ecmwf/anemoisecurity please ignore the review request, It was created by mistake during a merge with main.

Note to self: "do not use VSCode git UI ever again" :)

[JPXKQX]

Brilliant! Thanks for the great contribution Francesco!

[matschreiner]

Hi @frazane — this looks awesome, thanks a lot for implementing it!
I was on vacation last week so I didn’t get a chance to review earlier.
Would it make sense for me to open an issue with a few comments?

1. When zooming in, navigation becomes a tricky because the rotation sensitivity doesn’t seem aligned with the zoom level. Even a small drag rotates the globe quite a lot when zoomed in.

2. Could node scaling depending on connection be optional? In some graphs - for example on stretched grid the global nodes become extremely small compared to the cutout nodes e.g., in the screenshot below there are actually four global nodes, but they’re too tiny to see.

3. Why are the nodes rendered as triangles instead of circles/dots?

[frazane]

Hi @matschreiner,

I think it makes more sense to just directly open PRs with the proposed improvements. Both of the first two points seem sensible to me! And to answer to your question, the nodes are technically spheres:

https://github.com/ecmwf/anemoi-core/pull/609/files#diff-c18aaba70533311a3b8722c6808aa4fa4d71f28ec153693659fbc307dbb3c47fR340

just with a very small number of vertices so the rendering is faster.