:sparkles: add tests (13050dc2) · Commits · 钟慕尧 / deepchem

deepchem/feat/graph_data.py

+19 −0

Original line number	Diff line number	Diff line
		@@ -26,6 +26,13 @@ class GraphData:
		The number of edges in the graph
		num_edges_features: int, optional (default None)
		The number of features per edge in the graph

		Examples
		--------
		>>> import numpy as np
		>>> node_features = np.random.rand(5, 10)
		>>> edge_index = np.array([[0, 1, 2, 2, 3], [1, 2, 3, 3, 4]], dtype=np.int)
		>>> Graph(node_features=node_features, edge_index=edge_index)
		"""

		def __init__(
		@@ -114,6 +121,18 @@ class BatchGraphData(GraphData):
		----------
		graph_index: np.ndarray, dtype int
		This vector indicates which graph the node belongs with shape [num_nodes,]

		Examples
		--------
		>>> import numpy as np
		>>> node_features_list = np.random.rand(2, 5, 10)
		>>> edge_index_list = np.array([
		... [[0, 1, 2, 2, 3], [1, 2, 3, 3, 4]],
		... [[0, 1, 2, 2, 3], [1, 2, 3, 3, 4]],
		... ], dtype=np.int)
		>>> graphs = [Graph(node_features, edge_index) for node_features, edge_index
		... in zip(node_features_list, edge_index_list)]
		>>> BatchGraphData(graphs=graphs)
		"""

		def __init__(self, graphs: Sequence[GraphData]):

+10 −2

Original line number	Diff line number	Diff line
		@@ -8,8 +8,8 @@ from deepchem.utils.typing import PymatgenStructure
		from deepchem.feat import MaterialStructureFeaturizer
		from deepchem.feat.graph_data import GraphData

		# FIXME: it is better to add this json to DeepChem AWS
		ATOM_INIT_JSON_URL = 'https://raw.githubusercontent.com/txie-93/cgcnn/master/data/sample-regression/atom_init.json'

		ATOM_INIT_JSON_URL = 'https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/atom_init.json'


		class CGCNNFeaturizer(MaterialStructureFeaturizer):
		@@ -33,6 +33,14 @@ class CGCNNFeaturizer(MaterialStructureFeaturizer):
		----------
		.. [1] T. Xie and J. C. Grossman, Phys. Rev. Lett. 120, 2018.

		Examples
		--------
		>>> import pymatgen as mg
		>>> lattice = mg.Lattice.cubic(4.2)
		>>> structure = mg.Structure(lattice, ["Cs", "Cl"], [[0, 0, 0], [0.5, 0.5, 0.5]])
		>>> featurizer = CGCNNFeaturizer()
		>>> features = featurizer.featurize([structure])

		Note
		----
		This class requires Pymatgen to be installed.

+7 −0

Original line number	Diff line number	Diff line
		@@ -30,6 +30,13 @@ class ElementPropertyFingerprint(MaterialCompositionFeaturizer):
		.. [3] Matminer: Ward, L. et al. Comput. Mater. Sci. 152, 60-69 (2018).
		.. [4] Pymatgen: Ong, S.P. et al. Comput. Mater. Sci. 68, 314-319 (2013).

		Examples
		--------
		>>> import pymatgen as mg
		>>> comp = mg.Composition("Fe2O3")
		>>> featurizer = ElementPropertyFingerprint()
		>>> features = featurizer.featurize([comp])

		Note
		----
		This class requires matminer and Pymatgen to be installed.

+9 −1

Original line number	Diff line number	Diff line
		@@ -31,6 +31,14 @@ class SineCoulombMatrix(MaterialStructureFeaturizer):
		----------
		.. [1] Faber et al. Inter. J. Quantum Chem. 115, 16, 2015.

		Examples
		--------
		>>> import pymatgen as mg
		>>> lattice = mg.Lattice.cubic(4.2)
		>>> structure = mg.Structure(lattice, ["Cs", "Cl"], [[0, 0, 0], [0.5, 0.5, 0.5]])
		>>> featurizer = SineCoulombMatrix(max_atom=2)
		>>> features = featurizer.featurize([structure])

		Note
		----
		This class requires matminer and Pymatgen to be installed.
		@@ -47,7 +55,7 @@ class SineCoulombMatrix(MaterialStructureFeaturizer):
		Return flattened vector of matrix eigenvalues.
		"""

		self.max_atoms = int(max_atoms)
		self.max_atoms = max_atoms
		self.flatten = flatten

		def _featurize(self, struct: PymatgenStructure) -> np.ndarray:

+1 −1

Original line number	Diff line number	Diff line
		@@ -90,6 +90,6 @@ class TestGraph(unittest.TestCase):

		# check to_pyg_data function
		targets = np.array([1, 2, 3], dtype=np.float)
		batch = BatchGraphData.to_pyg_data(graph_list=graphs, targets=targets)
		batch = BatchGraphData.to_pyg_data(graphs=graphs, targets=targets)
		from torch_geometric.data import Batch
		assert isinstance(pyg_graph, Batch)