a-r-j / graphein

Protein Graph Library
https://graphein.ai/
MIT License
1.01k stars 126 forks source link

KeyError: 'kind' for split edge index by edge kind when converting nx to pyg #280

Closed rg314 closed 1 year ago

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

Originally posted by @rg314 in https://github.com/a-r-j/graphein/issues/234#issuecomment-1467928043

a-r-j commented 1 year ago

Good catch! i think all edges added through Graphein's methods will specify a kind. This would not necessarily be true for user-defined edge functions though so this should be fixed.

rg314 commented 1 year ago

I suggest:

        # Split edge index by edge kind
        kind_strs = np.array(list(map(lambda x: "_".join(x), data.get("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 and data.get('kind'):
            del data["kind"]

I'll PR

anton-bushuiev commented 1 year ago

Hi, @rg314! Yes, it's my mistake. I was considering kind mandatory. Great that you pointed it out.

rg314 commented 1 year ago

Good catch! i think all edges added through Graphein's methods will specify a kind. This would not necessarily be true for user-defined edge functions though so this should be fixed.

I'm not sure if the func graphein.rna.graphs.construct_graph has kind