Filter graph based on indices

[ibiscp]

I have a graph and I would like to sample it based on the indices of a given node.

For example, let's say that I have the following graph and I would like to select 'paper': torch.tensor([3, 4]), but I would like to receive a full graph back containing also author, conference, and all edges that are related to these paper indices.

HeteroData(
    paper={ x=[10, 16] },
    author={ x=[5, 32] },
    conference={ x=[5, 8] },
    (paper, cites, paper)={ edge_index=[2, 50] },
    (author, to, paper)={ edge_index=[2, 30] },
    (paper, to, conference)={ edge_index=[2, 25] }
)

I tried to use subgraph function, but I could not find an easy way to solve it.

Does anyone have any idea that can lead me to the right path?

[rusty1s]

data = data.subgraph({
    'paper': torch.tensor([3, 4]),
    'author': torch.arange(data['author'].num_nodes,
    'conference': torch.arange(data['conference'].num_nodes,
})

[ibiscp]

Thank you very much!

But it still does not solve my problem, let me present you with another example.

Let's start with the following graph:

num_authors = 2
num_authors_features = 32
num_papers = 4
num_paper_features = 16
num_conference = 8
num_conference_features = 8

data = HeteroData()

data['author'].x = torch.randn(num_authors, num_authors_features)
data['paper'].x = torch.randn(num_papers, num_paper_features)
data['conference'].x = torch.randn(num_conference, num_conference_features)

data['author', 'writes', 'paper'].edge_index = torch.vstack([torch.tensor([0,0,1,1]), torch.tensor([0,1,2,3])])
data['paper', 'to', 'conference'].edge_index = torch.vstack([torch.tensor([0,0,1,1,2,2,3,3]), torch.tensor([0,1,2,3,4,5,6,7])])

The output graph is

HeteroData(
  author={ x=[2, 32] },
  paper={ x=[4, 16] },
  conference={ x=[8, 8] },
  (author, writes, paper)={ edge_index=[2, 4] },
  (paper, to, conference)={ edge_index=[2, 8] }
)

Now, I want to filter my graph and select everything related to author 0

author_0 = data.subgraph({
    'author': torch.tensor([0]),
    'paper': torch.arange(data['author'].num_nodes),
    'conference': torch.arange(data['conference'].num_nodes),
})

Output

HeteroData(
  author={ x=[1, 32] },
  paper={ x=[2, 16] },
  conference={ x=[8, 8] },
  (author, writes, paper)={ edge_index=[2, 2] },
  (paper, to, conference)={ edge_index=[2, 4] }
)

It removed author 1, but the rest of the graph is still there.

Another example, now let's say I want the graph for conference 0

conference_0 = data.subgraph({
    'author': torch.arange(data['author'].num_nodes),
    'paper': torch.arange(data['author'].num_nodes),
    'conference': torch.tensor([0]),
})

HeteroData(
  author={ x=[2, 32] },
  paper={ x=[2, 16] },
  conference={ x=[1, 8] },
  (author, writes, paper)={ edge_index=[2, 2] },
  (paper, to, conference)={ edge_index=[2, 1] }
)

And the same happens for paper

Is it possible to receive only nodes directly linked to the requested nodes?

[ibiscp]

I have a graph constructed from a big dataset, to simplify, imagine that I have 100K users, 1M+ transactions, and some other entities such as transaction categories and subcategories. I'm training a GNN on CPU and I would like to train on a small sample of this data (10K users), to obtain some embeddings. I could filter the data before constructing the graph, but I imagined that with the subgraph function I could get this kind of sample, a subgraph based on the nodes I selected, without isolated nodes and disconnected links. What would be your approach to solving this problem?
For me, it makes sense that a subgraph function could have a cascade parameter, but I'm not sure if I have the right approach here, or if it makes sense.

[rusty1s]

The recommended approach for this would be the usage of k_hop_subgraph or creating mini-batches via NeighborLoader. Does that work for you?

[LukeLIN-web]

I have a graph constructed from a big dataset, to simplify, imagine that I have 100K users, 1M+ transactions, and some other entities such as transaction categories and subcategories. I'm training a GNN on CPU and I would like to train on a small sample of this data (10K users), to obtain some embeddings. I could filter the data before constructing the graph, but I imagined that with the subgraph function I could get this kind of sample, a subgraph based on the nodes I selected, without isolated nodes and disconnected links. What would be your approach to solving this problem? For me, it makes sense that a subgraph function could have a cascade parameter, but I'm not sure if I have the right approach here, or if it makes sense.

I am sorry to ask an unrelated problem, what software do you use to draw this graph?

[EdisonLeeeee]

@LukeLIN-web I guess it was networkx.

[ibiscp]

@rusty1s Ok, I'll try that, thank you!
@EdisonLeeeee Yes, I used networkx