Merge pull request #126 from rbharath/vs-utils-scripts

"""

Transform vs-datasets into standard-form CSV files.

"""

from __future__ import print_function

from __future__ import division

from __future__ import absolute_import

import gzip

import argparse

import os

import csv

from rdkit import Chem

import cPickle as pickle

import warnings

from multiprocessing import Pool

import itertools

from functools import partial

import pandas as pd

def merge_dicts(*dict_args):

  """

  Given any number of dicts, shallow copy and merge into a new dict,

  precedence goes to key value pairs in latter dicts.

  """

  result = {}

  for dictionary in dict_args:

    result.update(dictionary)

  return result

def parse_args(input_args=None):

  parser = argparse.ArgumentParser()

  parser.add_argument(

      "--data-dir", required=1,

      help="Directory with data to be loaded.")

  parser.add_argument(

      "--out", required=1,

      help="Location to write output csv file.")

  parser.add_argument(

      "--num-cores", default=0, type=int,

      help="Number of cores to use for multiprocessing.")

  parser.add_argument(

      "--id-prefix", default="CID",

      help="Location to write output csv file.")

  parser.add_argument(

      "--overwrite", action="store_true",

      help="Overwrite partially processed files.")

  return parser.parse_args(input_args)

def load_shard(shard, shards_dir, id_prefix):

  if "sdf.gz" not in shard:

    return  

  print("Processing shard %s" % shard)

  shard = os.path.join(shards_dir, shard)

  with gzip.open(shard) as f:

    supp = Chem.ForwardSDMolSupplier(f)

    mols = [mol for mol in supp if mol is not None]

  mol_dict = mols_to_dict(mols, id_prefix)

  return mol_dict

def load_shards(shards_dir, id_prefix, worker_pool=None):

  all_mols = []

  shards = os.listdir(shards_dir)

  if worker_pool is None:

    for shard in shards:

      mol_dict = load_shard(shard, shards_dir=shards_dir, id_prefix=id_prefix)

      if mol_dict is not None: 

        all_mols.append(mol_dict)

  else:

    load_shard_partial = partial(

        load_shard, shards_dir=shards_dir, id_prefix=id_prefix)

    all_mols = worker_pool.map(load_shard_partial, shards)

  all_mols = [mol_dict for mol_dict in all_mols if mol_dict is not None]

  all_mols_dict = merge_dicts(*all_mols)

  return all_mols_dict

def mols_to_dict(mols, id_prefix, log_every_n=5000):

  """Turn list of rdkit mols to large dictionary."""

  print("About to process molecules")

  mol_dict = {}

  for ind, mol in enumerate(mols):

    if ind % log_every_n == 0:

      print("Processing mol %d" % ind)

    prop_names = mol.GetPropNames()

    CID_name = None

    for prop in prop_names:

      if "CID" in str(prop):

        CID_name = str(prop)

        break

    if CID_name is not None:

      mol_id = mol.GetProp(CID_name)

    else:

      # If mol_id is not set, then use isomeric smiles as unique identifier

      mol_id = Chem.MolToSmiles(mol, isomericSmiles=True)

    mol_dict[id_prefix + mol_id] = Chem.MolToSmiles(mol, isomericSmiles=True)

  return mol_dict

def get_target_names(targets_dir):

  """Read list of targets in collection from disk."""

  targets = [target for target in os.listdir(targets_dir)

             if "pkl.gz" in target]

  return [remove_extensions(target) for target in targets]

def process_target(target, targets_dir, overwrite):

  if "pkl.gz" not in target:

    return 

  print("Processing target %s" % target)

  target_file = os.path.join(targets_dir, target)

  with gzip.open(target_file) as f:

    df = pickle.load(f)

    csv_file = target_to_csv(

        targets_dir, df, os.path.basename(target), overwrite=overwrite)

  return csv_file

def process_targets(targets_dir, overwrite, worker_pool=None):

  csv_files = []

  targets = os.listdir(targets_dir)

  if worker_pool is None:

    for target in targets:

      csv_file = process_target(target, targets_dir, overwrite)

      if csv_file is not None:

        csv_files.append(csv_file)

  else:

    process_target_partial = partial(

        process_target, targets_dir=targets_dir, overwrite=overwrite)

    csv_files = worker_pool.map(process_target_partial, targets)

    csv_files = [csv_file for csv_file in csv_files if csv_file is not None]

  return csv_files

def remove_extensions(target_name):

  """Removes file extensions from given name"""

  target_name = os.path.basename(target_name)

  while "." in target_name:

    target_name = os.path.splitext(target_name)[0]

  return target_name

def target_to_csv(targets_dir, df, target_name, log_every_n=50000,

                  overwrite=False):

  """Converts the data in a target dataframe to a csv."""

  target = remove_extensions(target_name) 

  csv_file = os.path.join(targets_dir, target + ".csv")

  if not overwrite and os.path.isfile(csv_file):

    return csv_file

  target_names = get_target_names(targets_dir) 

  data_df = pd.DataFrame(columns=(["mol_id"] + target_names))

  data_df["mol_id"] = df["mol_id"]

  def get_outcome(row):

    if row["outcome"] == "active":

      return "1"

    elif row["outcome"] == "inactive":

      return "0"

    else:

      return "" 

  data_outcomes = df.apply(get_outcome, axis=1)

  for ind, outcome in enumerate(data_outcomes):

    data_df.set_value(ind, target, outcome)

    for other_target in target_names:

      if other_target != target:

        data_df.set_value(ind, other_target, "")

  #iterator = data_df.iterrows()

  data_df.fillna("")

  with open(csv_file, "wb") as f:

    data_df.to_csv(f)

  return csv_file

def join_dict_datapoints(old_record, new_record, target_names):

  """Merge two datapoints together."""

  assert old_record is not None

  # TODO(rbharath): BROKEN!

  assert new_record is not None

  assert old_record["mol_id"] == new_record["mol_id"]

  for target in target_names:

    if old_record[target] != "":

      continue 

    elif new_record[target] != "":

      old_record[target] = new_record[target]

  return old_record

def merge_mol_data_dicts(mol_dict, csv_files, target_names):

  """Merge data from target and molecule listings."""

  print("len(mol_dict) = %d" % len(mol_dict))

  num_missing = 0

  merged_data = {}

  fields = ["mol_id"] + target_names

  final_fields = ["mol_id", "smiles"] + target_names

  merged_df = pd.DataFrame(columns=final_fields)

  merge_pos, merge_map = 0, {}

  for ind, csv_file in enumerate(csv_files):

    print("Merging %d/%d targets" % (ind+1, len(csv_files)))

    data_df = pd.read_csv(csv_file, na_filter=False)

    data_df.fillna("")

    data_dicts = data_df.to_dict("records")

    for data_dict in data_dicts:

      # Trim unwanted indexing fields

      data_dict = {field: data_dict[field] for field in fields}

      mol_id = data_dict["mol_id"]

      if mol_id not in mol_dict:

        num_missing += 1

        continue

      mol_smiles = mol_dict[mol_id]

      data_dict["smiles"] = mol_smiles

      if mol_id not in merged_data:

        merged_data[mol_id] = data_dict

      else:

        merged_data[mol_id] = join_dict_datapoints(

            merged_data[mol_id], data_dict, target_names)

  print("Number of mismatched compounds: %d" % num_missing)

  return merged_data

def generate_csv(data_dir, id_prefix, out, overwrite, worker_pool=None):

  """Transforms a vs-dataset into a CSV file.

  Args:

    data_dir: Directory name. Should contain two subdirectories named

              "shards" and "targets". The "shards" directory should contain

              gzipped sdf files. The "targets" directory should contain

              one gzipped pkl file per molecular target.

    id_prefix: Desired prefix for compound names. Should be "TOX" or "CID".

    out: Name of csv file to write.

  """

  shards_dir = os.path.join(data_dir, "shards")

  targets_dir = os.path.join(data_dir, "targets")

  target_names = get_target_names(targets_dir)

  mol_dict = load_shards(shards_dir, id_prefix, worker_pool)

  csv_files = process_targets(targets_dir, overwrite, worker_pool)

  merged_dict = merge_mol_data_dicts(mol_dict, csv_files, target_names)

  merged_df = pd.DataFrame(merged_dict.values())

  merged_df.fillna("")

  with open(out, "wb") as f:

    print("Writing csv file to %s" % out)

    merged_df.to_csv(f, index=False)

def main():

  args = parse_args()

  data_dir = args.data_dir

  id_prefix = args.id_prefix

  num_cores = args.num_cores

  out = args.out

  overwrite = args.overwrite

  # Connect to running ipython server

  if num_cores > 0:

    p = Pool(processes=num_cores)

    generate_csv(data_dir, id_prefix, out, overwrite, worker_pool=p)

  else:

    generate_csv(data_dir, id_prefix, out, overwrite)

if __name__ == "__main__":

  main()