def test_atomic_number_preprocessor_with_tox21():
preprocessor = AtomicNumberPreprocessor()
dataset = SDFFileParser(preprocessor) \
.parse(get_tox21_filepath('train'))['dataset']
index = numpy.random.choice(len(dataset), None)
atoms, = dataset[index]
assert atoms.ndim == 1
assert atoms.dtype == numpy.int32
def test_atomic_number_preprocessor_with_tox21():
preprocessor = AtomicNumberPreprocessor()
# labels=None as default, and label information is not returned.
dataset = SDFFileParser(preprocessor).parse(get_tox21_filepath('train'))
index = numpy.random.choice(len(dataset), None)
atoms, = dataset[index]
assert atoms.ndim == 1 # (atom, )
assert atoms.dtype == numpy.int32
def test_rsgcn_preprocessor_with_tox21():
preprocessor = RSGCNPreprocessor()
# labels=None as default, and label information is not returned.
dataset = SDFFileParser(preprocessor)\
.parse(get_tox21_filepath('train'))['dataset']
index = numpy.random.choice(len(dataset), None)
atoms, adjacency = dataset[index]
assert atoms.ndim == 1 # (atom, )
assert atoms.dtype == numpy.int32
assert adjacency.ndim == 2
assert adjacency.dtype == numpy.float32
def test_nfp_preprocessor_with_tox21():
preprocessor = NFPPreprocessor()
dataset = SDFFileParser(preprocessor, postprocess_label=None).parse(
get_tox21_filepath('train'))
index = numpy.random.choice(len(dataset), None)
atoms, adjs = dataset[index]
assert atoms.ndim == 1 # (atom, )
assert atoms.dtype == numpy.int32
# (atom from, atom to)
assert adjs.ndim == 2
assert adjs.dtype == numpy.float32
def test_gat_preprocessor():
preprocessor = RelGATPreprocessor()
def postprocess_label(label_list):
# Set -1 to the place where the label is not found,
# this corresponds to not calculate loss with `sigmoid_cross_entropy`
return [-1 if label is None else label for label in label_list]
dataset = SDFFileParser(preprocessor, postprocess_label=postprocess_label
).parse(get_tox21_filepath('train'))["dataset"]
index = numpy.random.choice(len(dataset), None)
atoms, adjs = dataset[index]
assert atoms.ndim == 1 # (atom, )
assert atoms.dtype == numpy.int32
# (edge_type, atom from, atom to)
assert adjs.ndim == 3
assert adjs.dtype == numpy.float32
