Added saving capabilities (for sklearn) and got models to train and evaluate. Still buggy though.

from keras.layers.convolutional import Convolution3D, MaxPooling3D

from keras.utils import np_utils

from deep_chem.utils.preprocess import split_dataset

from deep_chem.utils.preprocess import tensor_dataset_to_numpy

from deep_chem.utils.evaluate import eval_model

from deep_chem.utils.evaluate import compute_r2_scores

from deep_chem.utils.load import get_target_names

from deep_chem.utils.load import process_datasets

from deep_chem.utils.evaluate import results_to_csv

from deep_chem.utils.save import save_model

from deep_chem.utils.save import load_model

from deep_chem.utils.evaluate import compute_model_performance

def parse_args(input_args=None):

  """Parse command-line arguments."""

                      help="Name of endpoint specifying train/test split.")

  featurize_cmd.add_argument("--smiles-endpoint", type=str, default="smiles",

                      help="Name of endpoint specifying SMILES for molecule.")

  featurize_cmd.add_argument("--id-endpoint", type=str, default=None,

                      help="Name of endpoint specifying unique identifier for molecule.\n"

                           "If none is specified, then smiles-endpoint is used as identifier.")

  featurize_cmd.add_argument("--threshold", type=float, default=None,

                      help="If specified, will be used to binarize real-valued prediction-endpoint.")

  featurize_cmd.add_argument("--name", required=1,

                      choices=["classification", "regression"],

                      help="Type of learning task.")

  group.add_argument("--input-transforms", nargs="+", default=[],

                      choices=["normalize", "truncate-outliers"],

                      choices=["normalize-and-truncate"],

                      help="Transforms to apply to input data.")

  group.add_argument("--output-transforms", nargs="+", default=[],

                      choices=["log", "normalize"],

                      help="Transforms to apply to output data.")

  group.add_argument("--feature-types", nargs="+", required=1,

                      help="Types of featurizations to use.")

                      help="Types of featurizations to use.\n"

                           "Each featurization must correspond to subdirectory in\n"

                           "generated data directory.")

  group.add_argument("--paths", nargs="+", required=1,

                      help="Paths to input datasets.")

  group.add_argument("--splittype", type=str, default="scaffold",

                      choices=["logistic", "rf_classifier", "rf_regressor",

                      "linear", "ridge", "lasso", "lasso_lars", "elastic_net",

                      "singletask_deep_network", "multitask_deep_network",

                      "3D_cnn"])

                      "3D_cnn"],

                      help="Type of model to build. Some models may allow for\n"

                           "further specification of hyperparameters. See flags below.")

  group = train_cmd.add_argument_group("Neural Net Parameters")

  group.add_argument("--n-hidden", type=int, default=500,

                      help="Number of hidden neurons for NN models.")

  group.add_argument("--learning-rate", type=float, default=0.01,

  group = train_cmd.add_argument_group("save")

  group.add_argument("--saved-out", type=str, required=1,

                  help="Location to save trained model.")

  train_cmd.set_defaults(func=train_model)

  eval_cmd = subparsers.add_parser("eval",

                help="Evaluate trained model on specified data.")

  eval_cmd.add_argument("--paths", nargs="+", required=1,

                      help="Paths to input datasets.")

  eval_cmd.add_argument("--splittype", type=str, default="scaffold",

  group = eval_cmd.add_argument_group("load model/data")

  group.add_argument("--saved-in", type=str, required=1,

                  help="Location from which to load saved model.")

  group.add_argument("--modeltype", required=1,

                      choices=["sklearn", "keras"],

                      help="Type of model to load.")

  # TODO(rbharath): Is there a way to get rid of this guy?

  group.add_argument("--mode", default="singletask",

                      choices=["singletask", "multitask"],

                      help="Type of model being built.")

  # TODO(rbharath): EXTREMELY AWKWARD!!! Both the train and evaluation have to

  # specify the set of input/output transforms desired. This seems like a major

  # API smell with many, many potentials for buginess. I think the right step

  # here is to add a new global sub-command "transform" which performs

  # data-transforms upon the input data to generate train/test splits.

  group = eval_cmd.add_argument_group("load-and-transform")

  group.add_argument("--task-type", default="classification",

                      choices=["classification", "regression"],

                      help="Type of learning task.")

  group.add_argument("--input-transforms", nargs="+", default=[],

                      choices=["normalize-and-truncate"],

                      help="Transforms to apply to input data.")

  group.add_argument("--output-transforms", nargs="+", default=[],

                      choices=["log", "normalize"],

                      help="Transforms to apply to output data.")

  group.add_argument("--feature-types", nargs="+", required=1,

                      help="Types of featurizations to use.\n"

                           "Each featurization must correspond to subdirectory in\n"

                           "generated data directory.")

  group.add_argument("--paths", nargs="+", required=1,

                      help="Paths to evaluation datasets.")

  # TODO(rbharath): There is something awkward here in that we shouldn't have

  # to specify a split to obtain the test-set right? But I'm not sure what the

  # better method is here sicne often the test-set isn't actually stratified

  # out. When we are doing featurization, should we actually do a hard split

  # and write train/test to separate locations? That might actually the more

  # elegant path. 

  group.add_argument("--splittype", type=str, default="scaffold",

                       choices=["scaffold", "random", "specified"],

                       help="Type of train/test data-splitting.\n"

                            "scaffold uses Bemis-Murcko scaffolds.\n"

                            "specified requires that split be in original data.")

  eval_cmd.add_argument("--compute-aucs", type=bool, default=False,

                            "specified requires that split be in original data.\n"

                            "Evaluation performed upon this split of specified data.")

  group = eval_cmd.add_argument_group("metrics")

  group.add_argument("--compute-aucs", action="store_true", default=False,

                      help="Compute AUC for trained models on test set.")

  eval_cmd.add_argument("--compute-r2s", type=bool, default=False,

  group.add_argument("--compute-r2s", action="store_true", default=False,

                     help="Compute R^2 for trained models on test set.")

  eval_cmd.add_argument("--compute-rms", type=bool, default=False,

  group.add_argument("--compute-rms", action="store_true", default=False,

                     help="Compute RMS for trained models on test set.")

  eval_cmd.add_argument("--csv-out", type=str, default=None,

  group.add_argument("--csv-out", type=str, default=None,

                     help="Outputted predictions on the test set.")

  eval_cmd.set_defaults(func=eval_trained_model)

  return parser.parse_args(input_args)

      args.field_types, args.prediction_endpoint, args.smiles_endpoint,

      args.threshold, args.delimiter)

  generate_targets(df, mols, args.prediction_endpoint, args.split_endpoint,

      args.smiles_endpoint, out_y_pkl, out_sdf)

  generate_features(df, args.feature_endpoints, args.smiles_endpoint, out_x_pkl)

      args.smiles_endpoint, args.id_endpoint, out_y_pkl, out_sdf)

  generate_features(df, args.feature_endpoints, args.smiles_endpoint,

                    args.id_endpoint, out_x_pkl)

  generate_fingerprints(args.name, args.out)

  generate_descriptors(args.name, args.out)

def train_model(args):

  """Builds model from featurized data."""

  paths = args.paths

  targets = get_target_names(paths)

  targets = get_target_names(args.paths)

  task_types = {target: args.task_type for target in targets}

  input_transforms = args.input_transforms 

  output_transforms = {target: args.output_transforms for target in targets}

  # TODO(rbharath): The datatype (vector vs. tensor) should be automatically

  # detected in dataset_to_numpy

  datatype = "tensor" if args.model == "3D_cnn" else "vector"

  per_task_data = process_datasets(paths,

      input_transforms, output_transforms, feature_types=args.feature_types, 

      prediction_endpoint=args.prediction_endpoint,

      split_endpoint=args.split_endpoint,

  per_task_data = process_datasets(args.paths,

      args.input_transforms, output_transforms, feature_types=args.feature_types, 

      splittype=args.splittype, weight_positives=args.weight_positives,

      datatype=datatype, mode=args.mode)

  # TODO(rbharath): Bundle training params into a training_param dict that's passed

  # down to these functions.

      mode=args.mode)

  if args.model == "singletask_deep_network":

    models = fit_singletask_mlp(per_task_data, task_types, n_hidden=args.n_hidden,

      learning_rate=args.learning_rate, dropout=args.dropout,

      validation_split=args.validation_split)

  elif args.model == "3D_cnn":

    models = fit_3D_convolution(train_data, test_data, task_types,

        axis_length=args.axis_length, nb_epoch=args.n_epochs,

        batch_size=args.batch_size)

        nb_epoch=args.n_epochs, batch_size=args.batch_size)

  else:

    models = fit_singletask_models(per_task_data, args.model, task_types)

  # TODO(rbharath): Save trained model.

  if args.model in ["singletask_deep_network", "multitask_deep_network", "3D_cnn"]:

    modeltype = "keras"

  else:

    modeltype = "sklearn"

  save_model(models, modeltype, args.saved_out)

def eval_trained_model(args):

  results, aucs, r2s, rms = compute_model_performance(per_task_data, models,

  model = load_model(args.modeltype, args.saved_in)

  targets = get_target_names(args.paths)

  task_types = {target: args.task_type for target in targets}

  output_transforms = {target: args.output_transforms for target in targets}

  per_task_data = process_datasets(args.paths,

      args.input_transforms, output_transforms, feature_types=args.feature_types, 

      splittype=args.splittype, weight_positives=False,

      mode=args.mode)

  results, aucs, r2s, rms = compute_model_performance(per_task_data, task_types, model, args.modeltype,

    args.compute_aucs, args.compute_r2s, args.compute_rms) 

  if args.csv_out is not None:

    results_to_csv(results, args.csv_out, task_type=args.task_type)

import numpy as np

import warnings

from deep_chem.utils.preprocess import dataset_to_numpy

from deep_chem.utils.preprocess import tensor_dataset_to_numpy

from deep_chem.utils.preprocess import labels_to_weights

from sklearn.metrics import mean_squared_error

from sklearn.metrics import roc_auc_score

from rdkit import Chem

from rdkit.Chem.Descriptors import ExactMolWt

def compute_model_performance(per_task_data, models, aucs=True, r2s=False, rms=False):

def compute_model_performance(per_task_data, task_types, models, modeltype,

    aucs=True, r2s=False, rms=False):

  """Computes statistics for model performance on test set."""

  all_results, auc_vals, r2s, rms = {}, {}, {}, {}

  all_results, auc_vals, r2_vals, rms_vals = {}, {}, {}, {}

  for index, target in enumerate(sorted(per_task_data.keys())):

    print "Evaluating model %d" % index

    print "Target %s" % target

    (train, _, _, _), (test, _, _, _) = per_task_data[target]

    model = models[target]

    results = eval_model(test, model, {target: task_types[target]}, 

                         # We run singletask models as special cases of

                         # multitask.

                         modeltype="keras_multitask")

                         modeltype=modeltype)

    all_results[target] = results[target]

    if aucs:

      auc_vals.update(compute_roc_auc_scores(results, task_types))

    if rms:

      rms_vals.update(compute_rms_scores(results, task_types))

  print "(aucs, r2s, rms)"

  print (aucs, r2s, rms)

  if aucs:

    print "Mean AUC: %f" % np.mean(np.array(auc_vals.values()))

  if r2s:

    Either sklearn, keras, or keras_multitask

  """

  # Extract features for test set and make preds

  if datatype == "vector":

  X, _, _ = dataset_to_numpy(test_set)

  elif datatype == "tensor":

    X, _, _ = tensor_dataset_to_numpy(test_set)

  if len(np.shape(X)) > 2:  # Dealing with 3D data

    if len(np.shape(X)) != 5:

      raise ValueError("Tensorial datatype must be of shape (n_samples, N, N, N, n_channels).")

    (n_samples, axis_length, _, _, n_channels) = np.shape(X)

    # TODO(rbharath): Modify the featurization so that it matches desired shaped. 

    X = np.reshape(X, (n_samples, axis_length, n_channels, axis_length, axis_length))

  else:

    raise ValueError("Datatype must be vector or tensor.")

  if modeltype == "keras_multitask":

    predictions = model.predict({"input": X})

    ypreds = []

  sdf = os.path.join(shards_dir, "%s-0.sdf.gz" % name)

  fingerprints = os.path.join(fingerprint_dir,

      "%s-fingerprints.pkl.gz" % name)

  # TODO(rbharath): There's a bit of ugliness here. featurize modifies the

  # smiles strings internally, which I suspect leads to some non-matching

  # smiles strings, hence dropping some compounds. featurize needs to be

  # modified so that it can take in lists of smiles rather than just sdf file.

  # FIXME: Make this directly call the CircularFingerprint featurizer in vs_utils.

  subprocess.call(["python", "-m", "vs_utils.scripts.featurize",

                   "--scaffolds", "--smiles",

                   sdf, fingerprints,

  elif field_type == "ndarray":

    return data 

def generate_targets(df, mols, prediction_endpoint, split_endpoint, smiles_endpoint, out_pkl, out_sdf):

def generate_targets(df, mols, prediction_endpoint, split_endpoint,

    smiles_endpoint, id_endpoint, out_pkl, out_sdf):

  """Process input data file, generate labels, i.e. y"""

  #TODO(enf, rbharath): Modify package unique identifier to take user-specified 

    #unique identifier instead of assuming smiles string

  labels_df = pd.DataFrame([])

  labels_df["mol_id"] = df[[id_endpoint]]

  labels_df["smiles"] = df[[smiles_endpoint]]

  labels_df["prediction"] = df[[prediction_endpoint]]

  if split_endpoint is not None:

    labels_df = df[[smiles_endpoint, prediction_endpoint, split_endpoint]]

  else:

    labels_df = df[[smiles_endpoint, prediction_endpoint]]

    labels_df["split"] = df[[split_endpoint]]

  # Write pkl.gz file

  with gzip.open(out_pkl, "wb") as f:

  scaffold = engine.get_scaffold(mol)

  return(scaffold)

def generate_features(df, feature_endpoints, smiles_endpoint, out_pkl):

def generate_features(df, feature_endpoints, smiles_endpoint, id_endpoint, out_pkl):

  if feature_endpoints is None:

    print("No feature endpoint specified by user.")

    return

  features_df["scaffolds"] = df[[smiles_endpoint]].apply(

    functools.partial(generate_scaffold, smiles_endpoint=smiles_endpoint),

    axis=1)

  features_df["mol_id"] = df[[smiles_endpoint]].apply(lambda s : "", axis=1)

  features_df["smiles"] = df[[smiles_endpoint]]

  features_df["mol_id"] = df[[id_endpoint]]

  with gzip.open(out_pkl, "wb") as f:

    pickle.dump(features_df, f, pickle.HIGHEST_PROTOCOL)

from deep_chem.utils.preprocess import transform_outputs

from deep_chem.utils.preprocess import transform_inputs

from deep_chem.utils.preprocess import dataset_to_numpy

from deep_chem.utils.preprocess import tensor_dataset_to_numpy

from deep_chem.utils.preprocess import multitask_to_singletask

from deep_chem.utils.preprocess import split_dataset

from deep_chem.utils.preprocess import to_arrays

from vs_utils.utils import ScaffoldGenerator

def process_datasets(paths, input_transforms, output_transforms,

    prediction_endpoint=None, split_endpoint=None, datatype="vector",

    feature_types=["fingerprints"], mode="multitask", splittype="random",

    seed=None, weight_positives=True):

    feature_types=["fingerprints"], mode="multitask",

    splittype="random", seed=None, weight_positives=True):

  """Extracts datasets and split into train/test.

  Returns a dict that maps target names to tuples.

    Seed used for random splits.

  """

  dataset = load_and_transform_dataset(paths, input_transforms, output_transforms,

      prediction_endpoint, split_endpoint=split_endpoint,

      feature_types=feature_types, weight_positives=weight_positives)

  arrays = {}

  if mode == "singletask":

      if len(data) == 0:

        continue

      train, test = split_dataset(dataset, splittype)

      train_data, test_data = to_arrays(train, test, datatype)

      train_data, test_data = to_arrays(train, test)

      arrays[target] = (train_data, test_data)

  elif mode == "multitask":

    train, test = split_dataset(dataset, splittype)

    train_data, test_data = to_arrays(train, test, datatype)

    train_data, test_data = to_arrays(train, test)

    arrays["all"] = (train_data, test_data)

  else:

    raise ValueError("Unsupported mode for process_datasets.")

        if "pkl.gz" in target_pickle]

  return target_names

def load_assays(paths, prediction_endpoint, split_endpoint=None, target_dir_name="targets"):

def load_assays(paths, target_dir_name="targets"):

  """Load regression dataset labels from assays.

  Returns a dictionary that maps smiles strings to label vectors.

      target_name = target_pickle.split(".")[0]

      with gzip.open(os.path.join(target_dir, target_pickle), "rb") as f:

        contents = pickle.load(f)

        if prediction_endpoint not in contents:

        if "prediction" not in contents:

          raise ValueError("Prediction Endpoint Missing.")

        for ind, smiles in enumerate(contents["smiles"]):

          measurement = contents[prediction_endpoint][ind]

          if split_endpoint is not None:

            splits[smiles] = contents[split_endpoint][ind]

          measurement = contents["prediction"][ind]

          if "split" is not None:

            splits[smiles] = contents["split"][ind]

          else:

            splits[smiles] = None

          # TODO(rbharath): There is some amount of duplicate collisions

          labels[smiles][target_name] = measurement 

  return labels, splits

def load_datasets(paths, prediction_endpoint, split_endpoint, target_dir_name="targets",

    feature_types=["fingerprints"]):

def load_datasets(paths, target_dir_name="targets", feature_types=["fingerprints"]):

  """Load both labels and fingerprints.

  Returns a dictionary that maps smiles to pairs of (fingerprint, labels)

  """

  data = {}

  molecules = load_molecules(paths, feature_types)

  labels, splits = load_assays(paths, prediction_endpoint, split_endpoint, target_dir_name)

  labels, splits = load_assays(paths, target_dir_name)

  for ind, smiles in enumerate(molecules):

    if smiles not in labels:

      continue

    assert np.isclose(pos_weight, neg_weight)

def load_and_transform_dataset(paths, input_transforms, output_transforms,

    prediction_endpoint, split_endpoint=None, labels_endpoint="labels", weight_positives=True,

    datatype="tensor", feature_types=["fingerprints"]):

    weight_positives=True, feature_types=["fingerprints"]):

  """Transform data labels as specified

  Parameters

    are performed in the order specified. An empty list corresponds to no

    transformations. Only for regression outputs.

  """

  dataset = load_datasets(paths, prediction_endpoint, split_endpoint,

      feature_types=feature_types)

  if datatype == "vector":

  dataset = load_datasets(paths, feature_types=feature_types)

  X, y, W = dataset_to_numpy(dataset, weight_positives=weight_positives)

  elif datatype == "tensor":

    X, y, W = tensor_dataset_to_numpy(dataset)

  y = transform_outputs(y, W, output_transforms,

      weight_positives=weight_positives)

  X = transform_inputs(X, input_transforms)

        labels[target] = -1

      else:

        labels[target] = y[s_index][t_index]

    datapoint[labels_endpoint] = labels

    datapoint["labels"] = labels

    datapoint["fingerprint"] = X[s_index]

    trans_data[smiles] = datapoint