Fix and improve the utilities for PyG

[anton-bushuiev]

Reference Issues/PRs

No Reference Issues/PRs

What does this implement/fix? Explain your changes

This PR fixes the bugs related to the processing of PyG data.

graphein.ml.conversion.convert_nx_to_pyg

1. Fix the coordinates bug. Currently, the function creates a data.coords tensor with an extra dimension because coords is also a “graph-level feature”.
2. Improve to convert everything that is possible to torch.Tensors. Makes further usage much more easier and a resulting data object much more PyG-like.
3. Extend to multiple edge_indices for multiple edge types.

graphein.ml.visualisation.plot_pyg_data

4. Fix the arguments for the nested plotly_protein_structure_graph. Currently, it lacks the positional value for node_size_feature and the order of the following arguments is completely broken.
5. Adapt coords processing to the change number 2 in convert_nx_to_pyg.

What testing did you do to verify the changes in this PR?

Pull Request Checklist

• Added a note about the modification or contribution to the ./CHANGELOG.md file (if applicable)
• Added appropriate unit test functions in the ./graphein/tests/* directories (if applicable)
• Modify documentation in the corresponding Jupyter Notebook under ./notebooks/ (if applicable)
• Ran python -m py.test tests/ and make sure that all unit tests pass (for small modifications, it might be sufficient to only run the specific test file, e.g., python -m py.test tests/protein/test_graphs.py)
• Checked for style issues by running black . and isort .

[a-r-j]

I've made some changes to the CI/CD in #244 which I expect will resolve the test failures here

[a-r-j]

It looks like we're running into an issue when it comes to collating the distance matrices into a batch.

I suppose we have three options:

1. flatten distance matrices: (e.g. n x n -> n^2 x 1)
2. add additional dimension: 1 x n x n which then becomes 1 x max(n) x max(n) where max(n) is the largest number of nodes in the batch and smaller proteins are padded appropriately
3. Drop it entirely since it's easy to compute in torch.

Do you have any strong opinions?

[anton-bushuiev]

That's a good question. I'll try it and let you know.

[anton-bushuiev]

Hi, @a-r-j!

From my perspective, (1) is much better than (2) but (3) is the best. I would not serialize distance matrices for several reasons:

• it's easy to compute them in torch, as you wrote
• they introduce additional O(n^2) memory overhead
• from my experience they are typically not needed after data preparation.

What do you think about it?

[a-r-j]

Very good points @anton-bushuiev.

I am not sure about not supporting it all. While distance matrices are easy to compute, we can still run into scenarios where there are matrix features we may want to include (e.g. Hbond map). Thus, I propose:

• Distance matrices should not be stored by default due to the memory overhead
• Users have to specifically include them in the columns arg if they are to be included
• If matrices are to be stored, they should be stored in a sparse format (e.g. using this.

[anton-bushuiev]

I agree with the first two points but I am not sure about sparse format. Working with protein graphs, distance matrices are always dense because physical laws allow only zeros on diagonals. That's why I think simple reshaping would be the best:

# Flatten before writing
data.dist_mat = data.dist_mat.reshape(data.num_nodes * data.num_nodes)
# Restore after reading
data.dist_mat = data.dist_mat.reshape((data.num_nodes, data.num_nodes))

Do I miss the scenarios when they may be sparse? What is Hbond map (can you please send some link to its usage in Graphein)?

[a-r-j]

LGTM, thanks for the contribution!

[sonarqubecloud]

Kudos, SonarCloud Quality Gate passed!   

0 Bugs
0 Vulnerabilities
0 Security Hotspots
0 Code Smells

No Coverage information
No Duplication information

[rg314]

Hi @anton-bushuiev,

Your commit for "Improve convert_nx_to_pyg" added some breaking changes. Edges do not necessarily have a kind unless I'm not understanding something?

graphein/graphein/ml/conversion.py

Line 307 in b308a58

# Split edge index by edge kind

        # Split edge index by edge kind
        kind_strs = np.array(list(map(lambda x: "_".join(x), data["kind"])))
        for kind in set(kind_strs):
            key = f"edge_index_{kind}"
            if key in self.columns:
                mask = kind_strs == kind
                data[key] = edge_index[:, mask]
        if "kind" not in self.columns:
            del data["kind"]