def test_crystal_4(data):
cg = CrystalGraphFeaturizer(atom_feature=lambda s: [1, 2, 3, 4], n_jobs=1)
tmp = cg.transform(data)
assert isinstance(tmp, pd.DataFrame)
assert tmp.shape == (2, 2)
assert tmp.values[0, 0][0].numpy().tolist() == [1, 2, 3, 4]
def test_crystal_3(data):
cg = CrystalGraphFeaturizer()
tmp = cg.transform(data)
assert isinstance(tmp, pd.DataFrame)
assert tmp.shape == (2, 2)
assert tmp.values[0, 0].shape == (16, 92)
assert tmp.values[0, 1][0].shape == (16, 12, 41)
assert tmp.values[0, 1][1].shape == (16, 12)
def test_crystal_2(data):
cg = CrystalGraphFeaturizer(max_num_nbr=15, radius=10, atom_feature='elements')
tmp = cg.node_features(data[0])
assert isinstance(tmp, torch.Tensor)
assert tmp.shape == (16, 58)
edges, ids = cg.edge_features(data[0])
assert isinstance(edges, torch.Tensor)
assert edges.shape == (16, 15, 51)
assert isinstance(ids, torch.Tensor)
assert ids.shape == (16, 15)
def test_crystal_1(data):
cg = CrystalGraphFeaturizer()
assert cg.feature_labels == ['atom_feature', 'bond_feature']
tmp = cg.node_features(data[0])
assert isinstance(tmp, torch.Tensor)
assert tmp.shape == (16, 92)
edges, ids = cg.edge_features(data[0])
assert isinstance(edges, torch.Tensor)
assert edges.shape == (16, 12, 41)
assert isinstance(ids, torch.Tensor)
assert ids.shape == (16, 12)