Fix and improve the utilities for PyG

anton-bushuiev commented 1 year ago

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`

Fix the coordinates bug. Currently, the function creates a data.coords tensor with an extra dimension because coords is also a “graph-level feature”.
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.
Extend to multiple edge_indices for multiple edge types.

`graphein.ml.visualisation.plot_pyg_data`

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.
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

[x] Added a note about the modification or contribution to the ./CHANGELOG.md file (if applicable)
[x] Added appropriate unit test functions in the ./graphein/tests/* directories (if applicable)
[ ] Modify documentation in the corresponding Jupyter Notebook under ./notebooks/ (if applicable)
[x] 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 commented 1 year ago

Taking another look at this. It looks like the switch from distance matrices of the shape N x N to 1 x N x N is breaking PyG style batching. Do you have any arguments against changing this back?

anton-bushuiev commented 1 year ago

No, I agree that N x N is the right option. However, I did not change this. It is caused by this line which was there before.

sonarcloud[bot] commented 1 year ago

Kudos, SonarCloud Quality Gate passed!

0 Bugs
0 Vulnerabilities
0 Security Hotspots
0 Code Smells

No Coverage information
No Duplication information

a-r-j commented 1 year ago

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

a-r-j commented 1 year ago

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:

flatten distance matrices: (e.g. n x n -> n^2 x 1)
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
Drop it entirely since it's easy to compute in torch.

Do you have any strong opinions?

anton-bushuiev commented 1 year ago

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

anton-bushuiev commented 1 year ago

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 commented 1 year ago

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 commented 1 year ago

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)?

sonarcloud[bot] commented 1 year ago

Kudos, SonarCloud Quality Gate passed!

0 Bugs
0 Vulnerabilities
0 Security Hotspots
0 Code Smells

No Coverage information
No Duplication information

rg314 commented 1 year ago

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?

https://github.com/a-r-j/graphein/blob/b308a58934c29e145666bf35a073b226c9b2b6d4/graphein/ml/conversion.py#L307

        # 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"]

a-r-j / graphein