mllam / neural-lam

Neural Weather Prediction for Limited Area Modeling
MIT License
64 stars 25 forks source link

Mesh Node Update: Paper vs. Code #11

Closed bet20ICL closed 4 months ago

bet20ICL commented 4 months ago

Hi, I wanted to ask about a small discrepancy I may have found between the paper and the code.

image

In the paper, the G2M GNN updates the mesh embeddings using the current mesh embedding, and the edge embeddings of any G2M edges.

However, in the code, the grid node embeddings are also passed into the InteractionNet, which concatenates the grid and mesh embeddings.

Am I misunderstanding anything here? If the discrepancy is intentional, could someone please give a brief justification why the change was done?

joeloskarsson commented 4 months ago

Hi, there's quite a lot going on here, I hope I can help clarify. Note that the InteractionNetwork should in general be viewed as performing all of eq. 2a - 2d, not just the mesh node update in 2d. For the G2M one we don't care about the updated edge embeddings, so 2b is not used.

Note the tilde on top of $\tilde{e}_{s \rightarrow r}^\text{G2M}$, meaning that the edge embedding updates from eq. 2a are the ones used here. And the grid embeddings $v_s^G$ are necessary to compute 2a. So that is why they go into the GNN in the code here.

bet20ICL commented 4 months ago

Thanks for the quick reply! It makes sense now, thanks!