<commit_before>#!/usr/bin/env python
"""
Parse assay data from PubChem BioAssay (PCBA).
"""

__author__ = "<NAME>"
__copyright__ = "Copyright 2014, Stanford University"
__license__ = "BSD 3-clause"

import argparse
import cPickle
import gzip
import numpy as np
import pandas as pd


def parse_args(input_args=None):
    """
    Parse command-line arguments.

    Parameters
    ----------
    input_args : list, optional
        Input arguments. If not provided, defaults to sys.argv[1:].
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', required=1,
                        help='Input data file.')
    parser.add_argument('-m', '--map', required=1,
                        help='Molecule ID to SMILES map.')
    parser.add_argument('-o', '--output', required=1,
                        help='Output filename.')
    subparsers = parser.add_subparsers(dest='subcommand')
    subparsers.add_parser('classification')
    regression = subparsers.add_parser('regression')
    regression.add_argument('columns', nargs='+',
                            help='Column indices to include.')
    return parser.parse_args(input_args)


def read_pcba_data(filename):
    """
    Read PCBA data.

    Parameters
    ----------
    filename : str
        PCBA data filename.
    """
    if filename.endswith('.gz'):
        with gzip.open(filename) as f:
            df = pd.read_csv(f)
    else:
        df = pd.read_csv(filename)

    # remove duplicates
    df = df.drop_duplicates('PUBCHEM_CID')

    return df


def read_map(filename):
    """
    Read ID->SMILES map.

    Parameters
    ----------
    filename : str
        ID->SMILES map filename.
    """
    if filename.endswith('.gz'):
        f = gzip.open(filename)
    else:
        f = open(filename)
    id_map = cPickle.load(f)
    f.close()
    return id_map


def map_smiles(df, id_map):
    """
    Get SMILES for PubChem CIDs.

    Parameters
    ----------
    df : DataFrame
        DataFrame containing PCBA data.
    id_map : dict
        ID->SMILES map.
    """
    smiles = []
    indices = []
    for i, cid in enumerate(df.PUBCHEM_CID):
        name = 'CID{}'.format(cid)  # no bare IDs in map
        try:
            smiles.append(id_map[name])
            indices.append(i)
        except KeyError:
            pass
    return np.asarray(smiles), np.asarray(indices)


def save(smiles, targets, filename):
    """
    Save SMILES and targets to disk.

    Parameters
    ----------
    smiles : array_like
        SMILES for molecules.<commit_msg>Convert CIDs to int<commit_after>#!/usr/bin/env python
"""
Parse assay data from PubChem BioAssay (PCBA).
"""

__author__ = "<NAME>"
__copyright__ = "Copyright 2014, Stanford University"
__license__ = "BSD 3-clause"

import argparse
import cPickle
import gzip
import numpy as np
import pandas as pd


def parse_args(input_args=None):
    """
    Parse command-line arguments.

    Parameters
    ----------
    input_args : list, optional
        Input arguments. If not provided, defaults to sys.argv[1:].
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', required=1,
                        help='Input data file.')
    parser.add_argument('-m', '--map', required=1,
                        help='Molecule ID to SMILES map.')
    parser.add_argument('-o', '--output', required=1,
                        help='Output filename.')
    subparsers = parser.add_subparsers(dest='subcommand')
    subparsers.add_parser('classification')
    regression = subparsers.add_parser('regression')
    regression.add_argument('columns', nargs='+',
                            help='Column indices to include.')
    return parser.parse_args(input_args)


def read_pcba_data(filename):
    """
    Read PCBA data.

    Parameters
    ----------
    filename : str
        PCBA data filename.
    """
    if filename.endswith('.gz'):
        with gzip.open(filename) as f:
            df = pd.read_csv(f)
    else:
        df = pd.read_csv(filename)

    # remove duplicates
    df = df.drop_duplicates('PUBCHEM_CID')

    return df


def read_map(filename):
    """
    Read ID->SMILES map.

    Parameters
    ----------
    filename : str
        ID->SMILES map filename.
    """
    if filename.endswith('.gz'):
        f = gzip.open(filename)
    else:
        f = open(filename)
    id_map = cPickle.load(f)
    f.close()
    return id_map


def map_smiles(df, id_map):
    """
    Get SMILES for PubChem CIDs.

    Parameters
    ----------
    df : DataFrame
        DataFrame containing PCBA data.
    id_map : dict
        ID->SMILES map.
    """
    smiles = []
    indices = []
    for i, cid in enumerate(df.PUBCHEM_CID):
        if np.isnan(cid):
            continue
        cid = int(cid)  # CIDs are often read as floats
        name = 'CID{}'.format(cid)  # no bare IDs in map
        if name in id_map:
            smiles.append(id_map[name])
            indices.append(i)
    return np.asarray(smiles), np.asarray(indices)


def save(smiles, targets, filename):
    """
    Save SMILES and targets to disk.

    Parameters
    ----------
    smiles : array_like
        SMILES for molecules.
{'fraction_non_alphanumeric': 0.13128976286871025, 'fraction_numerical': 0.003470213996529786, 'mean_word_length': 2.0825906120023765, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7473041', 'n_tokens_mistral': 1668, 'n_tokens_neox': 1602, 'n_words': 399}