def test_schnet_preprocessor_with_tox21(): preprocessor = SchNetPreprocessor() dataset = SDFFileParser(preprocessor, postprocess_label=None).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 # (atom from, atom to) assert adjs.ndim == 2 assert adjs.dtype == numpy.float32
def test_schnet_preprocessor_default(): preprocessor = SchNetPreprocessor() dataset = SmilesParser(preprocessor).parse( ['C#N', 'Cc1cnc(C=O)n1C', 'c1ccccc1'])['dataset'] 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_schnet_preprocessor_assert_raises(): with pytest.raises(ValueError): pp = SchNetPreprocessor(max_atoms=3, out_size=2) # NOQA
def pp(): return SchNetPreprocessor()