Message Passing Module #516
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
dario-coscia
added
enhancement
|
Code Coverage SummaryResults for commit: 1c6bef4 Minimum allowed coverage is ♻️ This comment has been updated with latest results |
GiovanniCanali
force-pushed
the
messagepassing
branch
from
b17fd16 to
23b3560
Compare
|
Hi @AleDinve @GiovanniCanali ! How is it going with this? |
Hi @dario-coscia, I need to fix some minor issues with InteractionNetwork, and then I will fix EGNN. Also, tests will be implemented. @AleDinve agreed to take care of the remaining classes. |
Yes, I confirm, I will have a tentative implementation of the classes assigned to me by the weekend. |
GiovanniCanali
force-pushed
the
messagepassing
branch
2 times, most recently
from
5727bcc to
426da3e
Compare
dario-coscia
force-pushed
the
messagepassing
branch
from
f38ccbd to
1c6bef4
Compare
dario-coscia
commented
Apr 25, 2025
There was a problem hiding this comment.
Very good @GiovanniCanali and @AleDinve ! I made few comments on the implementation of the various blocks.
I think we should think about inserting inside utils.py a simple function that checks integer types and values. For example (very minimalistic):
def check_values(value, positive=True, strict=True):
if positive and strict:
assert value >= 0
.....this would reduce a lot of lines of code inside the blocks.
| :return: The message to be passed. | ||
| :rtype: torch.Tensor | ||
| """ | ||
| return self.message_net(torch.cat((x_i, x_j), dim=-1)) |
There was a problem hiding this comment.
here the edge_attr is needed
| .. seealso:: | ||
|
|
||
| **Original reference** Köhler, J., Klein, L., & Noé, F. (2020, November). | ||
| Equivariant flows: exact likelihood generative learning for symmetric densities. |
|
|
||
| :param x: The node features. | ||
| :type x: torch.Tensor | LabelTensor | ||
| :param torch.Tensor edge_index: The edge indices. In the original formulation, |
There was a problem hiding this comment.
80 lines max, also LabelTensor is ok
| :return: The message to be passed. | ||
| :rtype: torch.Tensor | ||
| """ | ||
| r = torch.norm(x_i-x_j) |
There was a problem hiding this comment.
torch.norm(x_i - x_j) computes the global norm (a scalar), which is probably not what we want in message passing.
We want per-edge norms, use:
r = torch.norm(x_i - x_j, dim=1, keepdim=True)| the messages are aggregated from all nodes, not only from the neighbours. | ||
| :return: The updated node features. | ||
| :rtype: torch.Tensor | ||
| """ |
There was a problem hiding this comment.
We need to avoid self loops, as in the original paper:
from torch_geometric.utils import remove_self_loops
edge_index, _ = remove_self_loops(edge_index)|
|
||
| .. seealso:: | ||
|
|
||
| **Original reference** Schütt, K., Kindermans, P. J., Sauceda Felix, H. E., Chmiela, S., Tkatchenko, A., & Müller, K. R. (2017). |
| flow="source_to_target", | ||
| ): | ||
| """ | ||
| Initialization of the :class:`RadialFieldNetworkBlock` class. |
| :return: The updated node features. | ||
| :rtype: torch.Tensor | ||
| """ | ||
| return self.propagate(edge_index=edge_index, x=x, pos=pos) |
There was a problem hiding this comment.
same comment as above for self loops
| """ | ||
| Compute the message to be passed between nodes and edges. | ||
|
|
||
| :param x_j: Concatenation of the node position and the |
There was a problem hiding this comment.
missing doc for pos
| details. Default is "source_to_target". | ||
| :raises ValueError: If `node_feature_dim` is not a positive integer. | ||
| :raises ValueError: If `edge_feature_dim` is not a positive integer. | ||
| """ |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR fixes #515.
Here a tentative RoadMap
InteractionNetworkBlock: implements the interaction network (See sec.2 of https://arxiv.org/pdf/1704.01212)DeepTensorNetworkBlock: implements the Deep Tensor Neural Networks block (See sec.2 of https://arxiv.org/pdf/1704.01212)EGNNBlock: implements the E(n) Equivariant Graph Neural Network block(See sec.3 of https://arxiv.org/pdf/2102.09844)RadialFieldBlock: implements the radial field network block (See Tab1. of https://arxiv.org/pdf/2102.09844)SchnetBlock: implements the schnet block (See Tab1. of https://arxiv.org/pdf/2102.09844)Checklist