Initial MAT Featurizer Commit, WIP (8a3a86e2) · Commits · 钟慕尧 / deepchem

deepchem/feat/molecule_featurizers/mat_featurizer.py

+85 −71

Original line number	Diff line number	Diff line
		@@ -25,13 +25,17 @@ else:
		IntTensor = torch.cuda.IntTensor
		DoubleTensor = torch.cuda.DoubleTensor


		class MATFeaturizer(MolecularFeaturizer):

		def __init__(
		self,
		self.mol = mol,
		self.add_dummy_node = add_dummy_node,
		self.one_hot_formal_charge = one_hot_formal_charge,
		add_dummy_node=True,
		one_hot_formal_charge=True,
		):
		self.mol = mol
		self.add_dummy_node = add_dummy_node
		self.one_hot_formal_charge = one_hot_formal_charge

		def OHVector(value, array):
		if value not in array:
		@@ -40,7 +44,8 @@ class MATFeaturizer(MolecularFeaturizer):

		def atom_features(self, atom, one_hot_formal_charge=True):
		attrib = []
		attrib += OHVector(atom.getAtomicNum(), [5, 6, 7, 8, 9, 15, 16, 17, 35, 53, 999])
		attrib += OHVector(atom.getAtomicNum(),
		[5, 6, 7, 8, 9, 15, 16, 17, 35, 53, 999])
		attrib += OHVector(len(atom.GetNeighbors()), [0, 1, 2, 3, 4, 5])
		attrib += OHVector(atom.GetTotalNumHs(), [0, 1, 2, 3, 4])

		@@ -55,15 +60,21 @@ class MATFeaturizer(MolecularFeaturizer):
		return np.array(attrib, dtype=np.float32)

		def _featurize(self, mol, add_dummy_node, one_hot_formal_charge):
		node_features = np.array([atom_features(atom, one_hot_formal_charge) for atom in mol.getAtoms()])
		node_features = np.array(
		[atom_features(atom, one_hot_formal_charge) for atom in mol.getAtoms()])
		adjacency_matrix = np.eye(mol.GetNumAtoms())
		for bond in mol.GetBonds():
		starting_atom = bond.GetBeginAtom().GetIdx()
		ending_atom = bond.getEndAtom().GetIdx()
		adjacency_matrix[starting_atom, ending_atom] = adjacency_matrix[ending_atom, starting_atom] = 1
		adjacency_matrix[starting_atom, ending_atom] = adjacency_matrix[
		ending_atom, starting_atom] = 1

		conformer = mol.GetConformer()
		positional_matrix = np.array([[conformer.GetAtomPosition(k).x, conformer.GetAtomPosition(k).y, conformer.GetAtomPosition(k).z] for k in range(mol.GetNumAtoms())])
		positional_matrix = np.array([[
		conformer.GetAtomPosition(k).x,
		conformer.GetAtomPosition(k).y,
		conformer.GetAtomPosition(k).z
		] for k in range(mol.GetNumAtoms())])
		distance_matrix = pairwise_distances(positional_matrix)

		if add_dummy_node:
		@@ -72,11 +83,14 @@ class MATFeaturizer(MolecularFeaturizer):
		m[0, 0] = 1.0
		node_features = m

		m = np.zeros((adjacency_matrix.shape[0] + 1, adjacency_matrix.shape[1] + 1))
		m = np.zeros((adjacency_matrix.shape[0] + 1,
		adjacency_matrix.shape[1] + 1))
		m[1:, 1:] = adjacency_matrix
		adjacency_matrix = m

		m = np.full((distance_matrix.shape[0] + 1, distance_matrix.shape[1] + 1), 1e6)
		m = np.full((distance_matrix.shape[0] + 1, distance_matrix.shape[1] + 1),
		1e6)
		m[1:, 1:] = distance_matrix
		distance_matrix = m

		return node_features, adjacency_matrix, distance_matrix
		No newline at end of file

Admin message