Skip to content

Commit

Permalink
Yapf
Browse files Browse the repository at this point in the history
  • Loading branch information
Bharath Ramsundar authored and Bharath Ramsundar committed May 2, 2020
1 parent be59de1 commit b22fe8a
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 25 deletions.
35 changes: 27 additions & 8 deletions deepchem/utils/test/test_voxel_utils.py
View file
Expand Up @@ -4,6 +4,7 @@
from deepchem.utils import voxel_utils
from deepchem.utils import hash_utils


class TestVoxelUtils(unittest.TestCase):

def test_convert_atom_to_voxel(self):
Expand All @@ -12,7 +13,8 @@ def test_convert_atom_to_voxel(self):
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,)

Expand All @@ -24,7 +26,8 @@ def test_convert_pair_atom_to_voxel(self):
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,)
Expand All @@ -35,13 +38,21 @@ def test_voxelize_convert_atom(self):
atom_index = 2
box_width = 16
voxel_width = 1
voxels_per_edge = int(box_width/voxel_width)
voxels_per_edge = int(box_width / voxel_width)
get_voxels = voxel_utils.convert_atom_to_voxel
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
Expand All @@ -52,10 +63,18 @@ def test_voxelize_convert_atom_pair(self):
atom_index_pair = (2, 3)
box_width = 16
voxel_width = 1
voxels_per_edge = int(box_width/voxel_width)
voxels_per_edge = int(box_width / voxel_width)
get_voxels = voxel_utils.convert_atom_pair_to_voxel
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)
30 changes: 13 additions & 17 deletions deepchem/utils/voxel_utils.py
View file
Expand Up @@ -5,10 +5,8 @@
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
Expand Down Expand Up @@ -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,'
' coords xyz = %s, coords in box = %s' %
(atom_index, coordinates[atom_index], indices))
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.
Expand All @@ -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)


Expand Down Expand Up @@ -129,16 +127,14 @@ def voxelize(get_voxels,
voxels_per_edge, nb_channel),
"""
# Number of voxels per one edge of box to voxelize.
voxels_per_edge = int(box_width/voxel_width)
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():
Expand Down

0 comments on commit b22fe8a

Please sign in to comment.