Yapf (b22fe8aa) · Commits · 钟慕尧 / deepchem

deepchem/utils/test/test_voxel_utils.py

+27 −8

Original line number	Diff line number	Diff line
		@@ -4,6 +4,7 @@ import deepchem as dc
		from deepchem.utils import voxel_utils
		from deepchem.utils import hash_utils


		class TestVoxelUtils(unittest.TestCase):

		def test_convert_atom_to_voxel(self):
		@@ -12,7 +13,8 @@ class TestVoxelUtils(unittest.TestCase):
		atom_index = 2
		box_width = 16
		voxel_width = 1
		indices = voxel_utils.convert_atom_to_voxel(coordinates, atom_index, box_width, voxel_width)
		indices = voxel_utils.convert_atom_to_voxel(coordinates, atom_index,
		box_width, voxel_width)
		assert len(indices) == 1
		assert indices[0].shape == (3,)

		@@ -24,7 +26,8 @@ class TestVoxelUtils(unittest.TestCase):
		atom_index_pair = (2, 3)
		box_width = 16
		voxel_width = 1
		indices = voxel_utils.convert_atom_pair_to_voxel([coordinates1, coordinates2], atom_index_pair, box_width, voxel_width)
		indices = voxel_utils.convert_atom_pair_to_voxel(
		[coordinates1, coordinates2], atom_index_pair, box_width, voxel_width)
		assert len(indices) == 2
		assert indices[0].shape == (3,)
		assert indices[1].shape == (3,)
		@@ -40,8 +43,16 @@ class TestVoxelUtils(unittest.TestCase):
		hash_function = hash_utils.hash_ecfp
		feature_dict = {1: "C", 2: "CC"}
		nb_channel = 16
		features = voxel_utils.voxelize(get_voxels, box_width, voxel_width, hash_function, coordinates, feature_dict, nb_channel=nb_channel)
		assert features.shape == (voxels_per_edge, voxels_per_edge, voxels_per_edge, nb_channel)
		features = voxel_utils.voxelize(
		get_voxels,
		box_width,
		voxel_width,
		hash_function,
		coordinates,
		feature_dict,
		nb_channel=nb_channel)
		assert features.shape == (voxels_per_edge, voxels_per_edge, voxels_per_edge,
		nb_channel)

		def test_voxelize_convert_atom_pair(self):
		N = 5
		@@ -57,5 +68,13 @@ class TestVoxelUtils(unittest.TestCase):
		hash_function = hash_utils.hash_ecfp_pair
		feature_dict = {(1, 2): ("C", "O"), (2, 3): ("CC", "OH")}
		nb_channel = 16
		features = voxel_utils.voxelize(get_voxels, box_width, voxel_width, hash_function, coordinates, feature_dict, nb_channel=nb_channel)
		assert features.shape == (voxels_per_edge, voxels_per_edge, voxels_per_edge, nb_channel)
		features = voxel_utils.voxelize(
		get_voxels,
		box_width,
		voxel_width,
		hash_function,
		coordinates,
		feature_dict,
		nb_channel=nb_channel)
		assert features.shape == (voxels_per_edge, voxels_per_edge, voxels_per_edge,
		nb_channel)

deepchem/utils/voxel_utils.py

+13 −17

Original line number	Diff line number	Diff line
		@@ -5,10 +5,8 @@ import logging
		import numpy as np
		logger = logging.getLogger(__name__)

		def convert_atom_to_voxel(coordinates,
		atom_index,
		box_width,
		voxel_width):

		def convert_atom_to_voxel(coordinates, atom_index, box_width, voxel_width):
		"""Converts atom coordinates to an i,j,k grid index.

		This function offsets molecular atom coordinates by
		@@ -37,13 +35,13 @@ def convert_atom_to_voxel(coordinates,
		indices = np.floor(
		(coordinates[atom_index] + box_width / 2.0) / voxel_width).astype(int)
		if ((indices < 0) \| (indices >= box_width / voxel_width)).any():
		logger.warning(
		'Coordinates are outside of the box (atom id = %s,'
		logger.warning('Coordinates are outside of the box (atom id = %s,'
		' coords xyz = %s, coords in box = %s' %
		(atom_index, coordinates[atom_index], indices))

		return [indices]


		def convert_atom_pair_to_voxel(coordinates_tuple, atom_index_pair, box_width,
		voxel_width):
		"""Converts a pair of atoms to a list of i,j,k tuples.
		@@ -68,11 +66,11 @@ def convert_atom_pair_to_voxel(coordinates_tuple, atom_index_pair, box_width,

		indices_list = []
		indices_list.append(
		convert_atom_to_voxel(coordinates_tuple[0], atom_index_pair[0],
		box_width, voxel_width)[0])
		convert_atom_to_voxel(coordinates_tuple[0], atom_index_pair[0], box_width,
		voxel_width)[0])
		indices_list.append(
		convert_atom_to_voxel(coordinates_tuple[1], atom_index_pair[1],
		box_width, voxel_width)[0])
		convert_atom_to_voxel(coordinates_tuple[1], atom_index_pair[1], box_width,
		voxel_width)[0])
		return (indices_list)


		@@ -132,13 +130,11 @@ def voxelize(get_voxels,
		voxels_per_edge = int(box_width / voxel_width)
		if dtype == "np.int8":
		feature_tensor = np.zeros(
		(voxels_per_edge, voxels_per_edge, voxels_per_edge,
		nb_channel),
		(voxels_per_edge, voxels_per_edge, voxels_per_edge, nb_channel),
		dtype=np.int8)
		else:
		feature_tensor = np.zeros(
		(voxels_per_edge, voxels_per_edge, voxels_per_edge,
		nb_channel),
		(voxels_per_edge, voxels_per_edge, voxels_per_edge, nb_channel),
		dtype=np.float16)
		if feature_dict is not None:
		for key, features in feature_dict.items():

Admin message