nichecompass.data.node_level_split_mask

nichecompass.data.node_level_split_mask(data, val_ratio=0.1, test_ratio=0.0, split_key='x')

Split data on node-level into training, validation and test sets by adding node-level masks (train_mask, val_mask, test_mask) to the PyG Data object.

Parameters:

• data (Data) – PyG Data object to be split.

• val_ratio (float (default: 0.1)) – Ratio of nodes to be included in the validation split.

• test_ratio (float (default: 0.0)) – Ratio of nodes to be included in the test split.

• split_key (str (default: 'x')) – The attribute key of the PyG Data object that holds the ground truth labels. Only nodes in which the key is present will be split.

Return type:

Data

Returns:

data: PyG Data object with ´train_mask´, ´val_mask´ and ´test_mask´ attributes added.