bioverse.frameworks.pyg | Bioverse documentation

Skip to content

# Source code for bioverse.frameworks.pyg

    import awkward as ak
    import numpy as np
    import torch
    from torch_geometric.data import Data

    from ..framework import Framework

    class PygData(Data):

        def uncollate(self, y):
            if "sizes" in self:
                y = ak.unflatten(y, self.sizes, axis=0)
            y = ak.Array({"target": y})
            return y

class PygFramework(Framework): @classmethod def collate(cls, X, y=None, attr=[]): if X.resolution == "atom": num_vertices = X.toc["atom"].sum(axis=-1).sum(axis=-1).ravel() else: num_vertices = X.toc["residue"].sum(axis=-1).ravel() num_molecules = len(num_vertices) vertex2batch = torch.arange(num_molecules).repeat_interleave( ak.to_torch(num_vertices) ) if "molecule_graph" in X: offsets = np.insert(np.cumsum(num_vertices), 0, 0)[:-1] X.molecule_graph = X.molecule_graph + offsets return PygData( features=( ak.to_torch(X.vertex_features).float() if "vertex_features" in X else None ), token=(ak.to_torch(X.vertex_token).int() if "vertex_token" in X else None), pos=(ak.to_torch(X.vertex_pos).float() if "vertex_pos" in X else None), edge_index=( ak.to_torch(ak.concatenate(X.molecule_graph, axis=1)).long() if "molecule_graph" in X else None ), # num_nodes=ak.to_torch(len(X[i]["vertices"].features)), mask=(ak.to_torch(X.vertex_mask).bool() if "vertex_mask" in X else None), y=( ak.to_torch( ak.flatten(y["target"], axis=1) if "sizes" in y.fields else y["target"] ).float() if not y is None else None ), vertex2batch=vertex2batch, num_vertices=num_vertices, num_molecules=num_molecules, sizes=(y["sizes"] if not y is None and "sizes" in y.fields else None), **{attr: ak.to_torch(X.__getattr__(attr)) for attr in attr}, )