Commit f24cc8c6 authored by joegomes's avatar joegomes
Browse files

Fixed user_defined_features usage in featurize.py. Fixed usage of featurize...

Fixed user_defined_features usage in featurize.py.  Fixed usage of featurize function in pdbbind_nnscore.py
parent be5033d2
Loading
Loading
Loading
Loading
+2 −0
Original line number Diff line number Diff line
@@ -192,6 +192,8 @@ class DataFeaturizer(object):
    df["smiles"] = ori_df[[self.smiles_field]]
    for task in self.tasks:
      df[task] = ori_df[[task]]
    for feature in self.user_specified_features:
      df[feature] = ori_df[[feature]]
    if self.split_field is not None:
      df["split"] = ori_df[[self.split_field]]
    if self.protein_pdb_field is not None:
+4 −2
Original line number Diff line number Diff line
@@ -34,12 +34,14 @@ def featurize_and_split(input_file, feature_dir, samples_dir, train_dir,
  
  if feature_files is None:
    print("About to featurize.")
    feature_files = featurizer.featurize(input_file, feature_types,
                                         feature_dir, shard_size=8)
    samples = featurizer.featurize(input_file, feature_dir,
                                   samples_dir, shard_size=8)
    print("Completed Featurization")
  else:
    # Transform data into arrays for ML
    samples = FeaturizedSamples(samples_dir, feature_files,
                                reload_data=False)

  # Split into train/test
  train_samples, test_samples = samples.train_test_split(
      splittype, train_dir, test_dir)