<commit_before>"""ODDT pipeline framework for virtual screening"""
import csv
from os.path import isfile
from multiprocessing.dummy import Pool
from itertools import chain

from oddt import toolkit

def _parallel_helper(args):
    """Private helper to workaround Python 2 pickle limitations to paralelize methods"""
    obj, methodname, arg = args
    return getattr(obj, methodname)(**arg)

class virtualscreening:
    def __init__(self, n_cpu=-1, verbose=False):
        """Virtual Screening pipeline stack

        Parameters
        ----------
            n_cpu: int (default=-1)
                The number of parallel procesors to use

            verbose: bool (default=False)
                Verbosity flag for some methods
        """
        self._pipe = None
        self.n_cpu = n_cpu
        self.num_input = 0
        self.num_output = 0
        self.verbose = verbose
        # setup pool
        self._pool = Pool(n_cpu if n_cpu > 0 else None)

    def load_ligands(self, fmt, ligands_file, *args, **kwargs):
        """Loads file with ligands.

        Parameters
        ----------
            file_type: string
                Type of molecular file

            ligands_file: string
                Path to a file, which is loaded to pipeline

        """
        if fmt == 'mol2' and toolkit.backend == 'ob':
            if 'opt' in kwargs:
                kwargs['opt']['c'] = None
            else:
                kwargs['opt'] = {'c': None}
        new_pipe = self._ligand_pipe(toolkit.readfile(fmt, ligands_file, *args, **kwargs))
        self._pipe = chain(self._pipe, new_pipe) if self._pipe else new_pipe

    def _ligand_pipe(self, ligands):
        for mol in ligands:
            self.num_input += 1
            yield mol

    def apply_filter(self, expression, filter_type='expression', soft_fail = 0):
        """Filtering method, can use raw expressions (strings to be evaled in if statement, can use oddt.toolkit.Molecule methods, eg. 'mol.molwt < 500')
        Currently supported presets:
            * Lipinski Rule of 5 ('r5' or 'l5')
            * Fragment Rule of 3 ('r3')

        Parameters
        ----------
            expression: string or list of strings
                Expresion(s) to be used while filtering.

            filter_type: 'expression' or 'preset' (default='expression')
                Specify filter type: 'expression' or 'preset'. Default strings are treated as expressions.

            soft_fail: int (default=0)
                The number of faulures molecule can have to pass filter, aka. soft-fails.
        """
        if filter_type == 'expression':
            self._pipe = self._filter(self._pipe, expression, soft_fail = soft_fail)
        elif filter_type == 'preset':
            # define presets
            # TODO: move presets to another config file
            # Lipinski rule of 5's
            if expression.lower() in ['l5', 'ro5']:
                self._pipe = self._filter(self._pipe, ['mol.molwt < 500', 'mol.calcdesc(["HBA1"])["HBA1"] <= 10', 'mol.calcdesc(["HBD"])["HBD"] <= 5', 'mol.calcdesc(["logP"])["logP"] <= 5'], soft_fail = soft_fail)
            # Rule of three
            elif expression.lower() in ['ro3']:
                self._pipe = self._filter(self._pipe, ['mol.molwt < 300', 'mol.calcdesc(["HBA1"])["HBA1"] <= 3', 'mol.calcdesc(["HBD"])["HBD"] <= 3', 'mol.calcdesc(["logP"])["logP"] <= 3'], soft_fail = soft_fail)
            # PAINS filter
            elif expression.lower() in ['pains']:
                pains_smarts = {}
                with open(dirname(__file__)+'filter/pains.smarts') as pains_file:
                    csv_reader = csv.reader(pains_file, delimiter="\t")
                    for line in csv_reader:
                        if len(line) > 1:
                            pains_smarts[line[1][8:-2]] = line[0]
                self._pipe = self._filter_smarts(self._pipe, pains_smarts.values(), soft_fail = soft_fail)

    def _filter_smarts(self, pipe, smarts, soft_fail = 0):
        for mol in pipe:
            if type(smarts) is list:
                compiled_smarts = [toolkit.Smarts(s) for s in smarts]
                fail = 0
                for s in compiled_smarts:
                    if len(s.findall(mol)) > 0:
                        fail += 1
                    if fail > soft_fail:
                        break
                if fail <= soft_fail:
                    yield mol
            else:
                compiled_smarts = toolkit.Smarts(smarts)
                if len(compiled_smiles.findall(mol)) == 0:
                    yield mol

    def _filter(self, pipe, expression, soft_fail = 0):
        for mol in pipe:
            if type(expression) is list:<commit_msg>Fix filtering module<commit_after>"""ODDT pipeline framework for virtual screening"""
import csv
from os.path import dirname, isfile
from multiprocessing.dummy import Pool
from itertools import chain

from oddt import toolkit

def _parallel_helper(args):
    """Private helper to workaround Python 2 pickle limitations to paralelize methods"""
    obj, methodname, arg = args
    return getattr(obj, methodname)(**arg)

class virtualscreening:
    def __init__(self, n_cpu=-1, verbose=False):
        """Virtual Screening pipeline stack

        Parameters
        ----------
            n_cpu: int (default=-1)
                The number of parallel procesors to use

            verbose: bool (default=False)
                Verbosity flag for some methods
        """
        self._pipe = None
        self.n_cpu = n_cpu
        self.num_input = 0
        self.num_output = 0
        self.verbose = verbose
        # setup pool
        self._pool = Pool(n_cpu if n_cpu > 0 else None)

    def load_ligands(self, fmt, ligands_file, *args, **kwargs):
        """Loads file with ligands.

        Parameters
        ----------
            file_type: string
                Type of molecular file

            ligands_file: string
                Path to a file, which is loaded to pipeline

        """
        if fmt == 'mol2' and toolkit.backend == 'ob':
            if 'opt' in kwargs:
                kwargs['opt']['c'] = None
            else:
                kwargs['opt'] = {'c': None}
        new_pipe = self._ligand_pipe(toolkit.readfile(fmt, ligands_file, *args, **kwargs))
        self._pipe = chain(self._pipe, new_pipe) if self._pipe else new_pipe

    def _ligand_pipe(self, ligands):
        for mol in ligands:
            self.num_input += 1
            yield mol

    def apply_filter(self, expression, soft_fail = 0):
        """Filtering method, can use raw expressions (strings to be evaled in if statement, can use oddt.toolkit.Molecule methods, eg. 'mol.molwt < 500')
        Currently supported presets:
            * Lipinski Rule of 5 ('ro5' or 'l5')
            * Fragment Rule of 3 ('ro3')
            * PAINS filter ('pains')

        Parameters
        ----------
            expression: string or list of strings
                Expresion(s) to be used while filtering.

            soft_fail: int (default=0)
                The number of faulures molecule can have to pass filter, aka. soft-fails.
        """
        if expression in ['l5', 'ro5', 'ro3', 'pains']:
            # define presets
            # TODO: move presets to another config file
            # Lipinski rule of 5's
            if expression.lower() in ['l5', 'ro5']:
                self._pipe = self._filter(self._pipe, ['mol.molwt < 500', 'mol.calcdesc(["HBA1"])["HBA1"] <= 10', 'mol.calcdesc(["HBD"])["HBD"] <= 5', 'mol.calcdesc(["logP"])["logP"] <= 5'], soft_fail = soft_fail)
            # Rule of three
            elif expression.lower() in ['ro3']:
                self._pipe = self._filter(self._pipe, ['mol.molwt < 300', 'mol.calcdesc(["HBA1"])["HBA1"] <= 3', 'mol.calcdesc(["HBD"])["HBD"] <= 3', 'mol.calcdesc(["logP"])["logP"] <= 3'], soft_fail = soft_fail)
            # PAINS filter
            elif expression.lower() in ['pains']:
                pains_smarts = {}
                with open(dirname(__file__)+'/filter/pains.smarts') as pains_file:
                    csv_reader = csv.reader(pains_file, delimiter="\t")
                    for line in csv_reader:
                        if len(line) > 1:
                            pains_smarts[line[1][8:-2]] = line[0]
                self._pipe = self._filter_smarts(self._pipe, pains_smarts.values(), soft_fail = soft_fail)
        else:
            self._pipe = self._filter(self._pipe, expression, soft_fail = soft_fail)

    def _filter_smarts(self, pipe, smarts, soft_fail = 0):
        for mol in pipe:
            if type(smarts) is list:
                compiled_smarts = [toolkit.Smarts(s) for s in smarts]
                fail = 0
                for s in compiled_smarts:
                    if len(s.findall(mol)) > 0:
                        fail += 1
                    if fail > soft_fail:
                        break
                if fail <= soft_fail:
                    yield mol
            else:
                compiled_smarts = toolkit.Smarts(smarts)
                if len(compiled_smiles.findall(mol)) == 0:
                    yield mol

    def _filter(self, pipe, expression, soft_fail = 0):
        for mol in pipe:
            if type(expression) is list:
{'fraction_non_alphanumeric': 0.11467642941746528, 'fraction_numerical': 0.011090772046947346, 'mean_word_length': 1.8128407026044822, 'pattern_counts': {'":': 0, '<': 23, '<?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': '7453156', 'n_tokens_mistral': 2790, 'n_tokens_neox': 2579, 'n_words': 779}