Commit 542fafd0 authored by Nathan Frey's avatar Nathan Frey
Browse files

dimension fix

parent 698519d7

deepchem/feat/complex_featurizers/rdkit_grid_featurizer.py

+3 −2
Original line number Diff line number Diff line
@@ -449,6 +449,7 @@ class RdkitGridFeaturizer(ComplexFeaturizer): 
        else:

          features_dict[system_id] = np.concatenate(feature_arrays, axis=-1)



    # TODO(rbharath): Is this squeeze OK?

    features = np.squeeze(np.array(list(features_dict.values())))

    features = np.array(list(features_dict.values()))

    if self.nb_rotations == 0:  # squeeze out axis with dimension 1

      features = np.squeeze(features, axis=0)

    return features

deepchem/feat/tests/test_rdkit_grid_features.py

+2 −0
Original line number Diff line number Diff line
@@ -134,6 +134,8 @@ class TestRdkitGridFeaturizer(unittest.TestCase): 
  def test_rotations(self):

    featurizer = RdkitGridFeaturizer(

        nb_rotations=3,

        box_width=16.,

        voxel_width=1.,

        feature_types=['voxel_combined'],

        flatten=False,

        sanitize=True)

CRA Git | Maintained and supported by SUSTech CRA and CCSE