neutralize.py

#!/usr/bin/env python3

__author__ = 'NAME SURNAME'

import argparse
import sys
from read_input import read_input
from rdkit import Chem
from multiprocessing import Pool, cpu_count

pattern = Chem.MolFromSmarts("[+1!h0!$([*]~[-1,-2,-3,-4]),-1!$([*]~[+1,+2,+3,+4])]")

def neutralize_atoms(mol):
    # https://www.rdkit.org/docs/Cookbook.html#neutralizing-molecules
    mol = Chem.RemoveHs(mol)
    at_matches = mol.GetSubstructMatches(pattern)
    at_matches_list = [y[0] for y in at_matches]
    if len(at_matches_list) > 0:
        for at_idx in at_matches_list:
            atom = mol.GetAtomWithIdx(at_idx)
            chg = atom.GetFormalCharge()
            hcount = atom.GetTotalNumHs()
            atom.SetFormalCharge(0)
            atom.SetNumExplicitHs(hcount - chg)
            atom.UpdatePropertyCache()
    return mol


def neutralize_item(item):
    mol, mol_name = item
    try:
        mol = neutralize_atoms(mol)
        output = Chem.MolToSmiles(mol, isomericSmiles=True, canonical=True) + '\t' + mol_name + '\n'
    except:
        output = None
    return output


def neutralize(input_fname, output_fname, ncpu, verbose):

    pool = Pool(max(min(cpu_count(), ncpu), 1))

    input_format = 'smi' if input_fname is None else None

    fo = open(output_fname, "wt") if output_fname is not None else sys.stdout
    try:
        for i, line in enumerate(pool.imap(neutralize_item,
                                           read_input(input_fname, input_format=input_format),
                                           chunksize=1), 1):
            if line:
                fo.write(line)
            if verbose and i % 1000 == 0:
                sys.stderr.write(f'\rProcessed {i}')
        if verbose:
            sys.stderr.write(f'\n')
    finally:
        if output_fname is not None:
            fo.close()


def main():
    parser = argparse.ArgumentParser(description='Neutralize input structures. Explicit hydrogens will be removed')
    parser.add_argument('-i', '--input', metavar='FILENAME', required=False, default=None, type=str,
                        help='input SDF or SMILES file. If omitted STDIN will be read as SMILES.')
    parser.add_argument('-o', '--output', metavar='FILENAME', required=False, default=None, type=str,
                        help='output SMILES file. If omitted output will be redirected to STDOUT.')
    parser.add_argument('-c', '--ncpu', metavar='INTEGER', default=1, type=int,
                        help='number of cpu to use for calculation. Default: 1.')
    parser.add_argument('-v', '--verbose', action='store_true', default=False,
                        help='print progress to STDERR.')
    args = parser.parse_args()

    if args.input == "/dev/stdin":
        args.input = None

    neutralize(args.input, args.output, args.ncpu, args.verbose)


if __name__ == '__main__':
    main()