Error when running 'model_inference.ipynb'

[ChenXingLing]

I don't know how to fix this strange error

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[4], line 6
      3 pdb_files = os.listdir(pdb_dir)
      5 for pdb_file in tqdm(pdb_files):
----> 6     graph = get_processed_graph(pdb2graph(pdb_dir + pdb_file))
      7     collator = CollatorDiff()
      8     g_batch, ipa_batch = collator([graph])

File ~/MapDiff-1.0.1/data/generate_graph_cath.py:80, in get_processed_graph(data)
     70 def get_processed_graph(data):
     71     extra_x_feature = torch.cat([data.x[:, 20:], data.mu_r_norm], dim=1)
     72     graph = Data(
     73         x=data.x[:, :20],
     74         extra_x=extra_x_feature,
     75         pos=data.pos,
     76         atom_pos=data.atom_pos,
     77         edge_index=data.edge_index,
     78         edge_attr=data.edge_attr,
     79         ss=data.ss[:data.x.shape[0], :],
---> 80         sasa=data.x[:, 20]
     81     )
     82     return graph

IndexError: index 20 is out of bounds for dimension 1 with size 18

[Gabriel-QIN]

In MIF paper, it is said each residue is represented as a node in a graph. The edges are constructed using a k-nearest-neighbor (k-NN) approach:

Each node (residue) is connected to at most k neighbors that are within a cutoff distance of 30 Å.

If a residue has fewer than k neighbors within 30 Å, it is simply connected to all residues within that cutoff, so the number of edges per node is min(k, number of residues within 30 Å).

This ensures that the graph captures local structural context while remaining sparse, and it works for proteins of varying lengths.

So in short, the number of neighbors per node depends on the local density of residues rather than the total protein length.