def prot_graph_transform(item, atom_keys=['atoms'], label_key='scores'):
"""Transform for converting dataframes to Pytorch Geometric graphs, to be applied when defining a :mod:`Dataset <atom3d.datasets.datasets>`.
Operates on Dataset items, assumes that the item contains all keys specified in ``keys`` and ``labels`` arguments.
:param item: Dataset item to transform
:type item: dict
:param atom_keys: list of keys to transform, where each key contains a dataframe of atoms, defaults to ['atoms']
:type atom_keys: list, optional
:param label_key: name of key containing labels, defaults to ['scores']
:type label_key: str, optional
:return: Transformed Dataset item
:rtype: dict
"""
from torch_geometric.data import Data
import atom3d.util.graph as gr
for key in atom_keys:
node_feats, edge_index, edge_feats, pos = gr.prot_df_to_graph(
item[key])
item[key] = Data(node_feats,
edge_index,
edge_feats,
y=item[label_key],
pos=pos)
return item
def df_to_graph(struct_df, chain_res, label):
"""
Extracts atoms within 30A of CA atom and computes graph
"""
chain, resnum = chain_res
res_df = struct_df[(struct_df.chain == chain) & (struct_df.residue == resnum)]
if 'CA' not in res_df.name.tolist():
return None
ca_pos = res_df[res_df['name']=='CA'][['x', 'y', 'z']].astype(np.float32).to_numpy()[0]
kd_tree = scipy.spatial.KDTree(struct_df[['x','y','z']].to_numpy())
graph_pt_idx = kd_tree.query_ball_point(ca_pos, r=30.0, p=2.0)
graph_df = struct_df.iloc[graph_pt_idx].reset_index(drop=True)
ca_idx = np.where((graph_df.chain == chain) & (graph_df.residue == resnum) & (graph_df.name == 'CA'))[0]
if len(ca_idx) > 0:
return None
node_feats, edge_index, edge_feats, pos = gr.prot_df_to_graph(graph_df)
data = Data(node_feats, edge_index, edge_feats, y=label, pos=pos)
data.ca_idx = torch.LongTensor(ca_idx)
data.n_nodes = data.num_nodes
return data
def graph_transform(item, keys=['atoms']):
"""Transform for converting dataframes to graphs, to be applied when defining a :mod:`Dataset <atom3d.datasets.datasets>`.
Operates on Dataset items, assumes that the item contains a "labels" key and all keys specified in ``keys`` argument.
:param item: Dataset item to transform
:type item: dict
:param keys: list of keys to transform, where each key contains a dataset of atoms, defaults to ['atoms']
:type keys: list, optional
:return: Transformed Dataset item
:rtype: dict
"""
from torch_geometric.data import Data
import atom3d.util.graph as gr
for key in keys:
node_feats, edge_index, edge_feats, pos = gr.prot_df_to_graph(
item[key])
item[key] = Data(node_feats,
edge_index,
edge_feats,
y=item['label'],
pos=pos)
return item