kinoml.datasets.torch_geometric_datasets

Helper classes to convert between DatasetProvider objects and Dataset-like objects native to the PyTorch ecosystem

Module Contents

class kinoml.datasets.torch_geometric_datasets.AwkwardArrayGeometricDataset(data)

Bases: torch.utils.data.Dataset

Loads an Awkward array of Records suitable for PyTorch Geometric. It assumes the following:

• The Awkward array contains three fields: 0, 1 and 2

• 0: Conn. matrix –> Data’s edge_index

• 1: Node features –> Data’s x

• 2: y labels

If more attributes are needed, you need to modify __getitem__ logic

__len__()

__getitem__(index)

abstract __iter__()

__repr__()

__str__()

classmethod from_parquet(path, **kwargs)