DP5.py

import numpy as np
import qml
from qml.fchl import get_atomic_kernels
from scipy.stats import gaussian_kde as kde
import pickle
from scipy.stats import gmean
from pathlib import Path
from scipy import stats
import os
import pathos.multiprocessing as mp
import copy
import gzip

try:
    from openbabel.openbabel import OBConversion, OBMol, OBAtomAtomIter, OBMolAtomIter
except ImportError:
    from openbabel import *


c_distance = 4.532297920317418

class DP5data:

    def __init__(self,ScriptPath,Atoms):

        self.Atom_number = Atoms

        self.Cshifts = []  # Carbon shifts used in DP4 calculation
        self.Cexp = []  # Carbon experimental shifts used in DP4 calculation
        self.Clabels = []  # Carbon atom labels
        self.Cinds = []
        self.Hshifts = []  # Proton shifts used in DP4 calculation
        self.Hexp = []  # Proton experimental shifts used in DP4 calculation
        self.Hlabels = []  # Proton atom labels
        self.Cscaled = []  # Internally scaled carbon shifts
        self.Hscaled = []  # Internally scaled proton shifts

        self.ConfCshifts = []

        self.Compounds = [] #qml compound objects for each isomer
        self.AtomReps = [] # FCHL representations ordered by Clabels

        self.ScaledAtomProbs = [] #per atom dp5 scaled probabilities for all conformers

        self.CScaledprobs = []  # DP5 for isomers based on Carbon data

        self.BScaledAtomProbs = [] #per atom dp5 scaled probabilities boltzmann weighted

        self.DP5scaledprobs = []  # Final DP5S

        self.output = str()  # final DP4 output

        self.folded_scaled_errors = pickle.load(open(ScriptPath / "folded_scaled_errors.p", "rb"))

        #self.folded_unscaled_errors = pickle.load(open(ScriptPath / "folded_unscaled_errors.p", "rb"))

        if self.Atom_number < 86:

            with gzip.open(ScriptPath / "atomic_reps.gz", "rb") as f:

                self.atomic_reps = pickle.load(f)

        else:

            with gzip.open(ScriptPath / "frag_reps.gz", "rb") as f:

                self.atomic_reps = pickle.load(f)

        self.mean_abs_error = np.mean(abs(self.folded_scaled_errors))

        self.output = ""

def kde_probs(Isomers,dp5Data,sigma):

    def kde_probfunction(conf_shifts, conf_reps):

        scaled_probs = []

        errors = [abs(shift - exp) for shift, exp in zip(conf_shifts, dp5Data.Cexp[iso])]

        scaled_shifts = ScaleNMR(conf_shifts, dp5Data.Cexp[iso])

        scaled_errors = [abs(shift - exp) for shift, exp in zip(scaled_shifts, dp5Data.Cexp[iso])]

        for e, s_e, r in zip(errors, scaled_errors, conf_reps):

            # calculate similarites between this atom and those in the atomic representation test set

            K_sim = get_atomic_kernels(np.array([r]), dp5Data.atomic_reps, [sigma],
                                       cut_distance=c_distance)[0][0]

            K_sim = np.hstack((K_sim, K_sim))

            # calculate kde using K_sim as the weighting function

            if np.sum(K_sim) == 0:

                scaled_kde_estimator = kde(dp5Data.folded_scaled_errors)

            else:

                scaled_kde_estimator = kde(dp5Data.folded_scaled_errors, weights=K_sim)

            s_e_diff = abs(s_e - dp5Data.mean_abs_error)

            s_p = scaled_kde_estimator.integrate_box_1d(dp5Data.mean_abs_error - s_e_diff, dp5Data.mean_abs_error + s_e_diff)

            scaled_probs.append(s_p)

        return  scaled_probs

    #for each atom in the molecule calculate the atomic worry factor

    dp5Data.ScaledAtomProbs = [[] for i in range(len(Isomers))]


    for iso in range(len(Isomers)):

        res = [[] for i in dp5Data.AtomReps[iso]]

        dp5Data.ScaledAtomProbs[iso] = [[] for i in dp5Data.AtomReps[iso]]

        maxproc = 6

        pool = mp.Pool(maxproc)

        ind1 = 0

        for conf_shifts , conf_reps in zip(dp5Data.ConfCshifts[iso],dp5Data.AtomReps[iso] ) :


            res[ind1] = pool.apply_async(kde_probfunction,
                                         [conf_shifts,conf_reps])

            ind1 += 1

        for ind1 in range(len(res)):

            dp5Data.ScaledAtomProbs[iso][ind1] = res[ind1].get()

    return dp5Data


def ProcessIsomers(dp5Data, Isomers,Settings):


    OutputFolder = Path(Settings.OutputFolder)

    # extract calculated and experimental shifts and add to dp5Data instance

    # Carbon

    # make sure any shifts with missing peaks are removed from all isomers

    removedC = []

    no_exp_data = 0

    for iso in Isomers:

        if len(iso.Cexp) == 0:

            no_exp_data +=1

    if no_exp_data == len(Isomers):

        print("no experimental NMR data provided... quitting")

        quit()


    for iso in Isomers:

        dp5Data.Cexp.append([])
        dp5Data.Cshifts.append([])
        dp5Data.Clabels.append([])
        dp5Data.Cinds.append([])

        dp5Data.ConfCshifts.append([])

        a_ind = 0

        exp_inds = []


        for shift, exp, label in zip(iso.Cshifts, iso.Cexp, iso.Clabels):

            if exp != '':

                dp5Data.Cshifts[-1].append(shift)
                dp5Data.Cexp[-1].append(exp)
                dp5Data.Clabels[-1].append(label)
                dp5Data.Cinds[-1].append(int(label[1:]) - 1)
                exp_inds.append(a_ind)

            a_ind += 1

        if len(iso.ConformerCShifts) > 0:

            for conf_shifts in iso.ConformerCShifts:

                dp5Data.ConfCshifts[-1].append( [conf_shifts[e] for e in exp_inds])


    #write qml compound objects and atomic representations

    #check the number of atoms in the structures

    #if there are less than 86 (max number of atoms in a molecule in the training set) atoms

    if dp5Data.Atom_number < 86:

        for iso in Isomers:

            #open new xyz file

            InputFile = Path(iso.InputFile)

            #find conformer with the lowest energy

            dp5Data.AtomReps.append([])

            for i,geom in enumerate(iso.DFTConformers):

                xyz_file = open(str(OutputFolder / "dp5" /InputFile.stem) + "_" +str(i).zfill(3) + ".xyz", "w")

                xyz_file.write(str(len(iso.Atoms)) + "\n" + "\n")

                for atom, coords in zip(iso.Atoms, geom):

                    xyz_file.write(atom + " " + str(coords[0]) + " " + str(coords[1]) + " " + str(coords[2]) + "\n")

                xyz_file.close()

                dp5Data.Compounds.append(qml.Compound(xyz = str(Settings.OutputFolder/"dp5"/ InputFile.stem) +"_"+ str(i).zfill(3) + ".xyz"))

                dp5Data.Compounds[-1].generate_fchl_representation(max_size=86, cut_distance=c_distance)

                dp5Data.AtomReps[-1].append([])

                for C_l in iso.Clabels:

                    ind = int(C_l.split("C")[1])

                    dp5Data.AtomReps[-1][-1].append(dp5Data.Compounds[-1].representation[ind])

    #otherwise we need to fragment the molecule to radius of 3

    else:

        for iso in Isomers:

            #open new xyz file

            InputFile = Path(iso.InputFile)

            #find conformer with the lowest energy

            dp5Data.AtomReps.append([])

            for i,geom in enumerate(iso.DFTConformers):

                xyz_file = open(str(OutputFolder / "dp5" /InputFile.stem) + "_" +str(i).zfill(3) + ".xyz", "w")

                xyz_file.write(str(len(iso.Atoms)) + "\n" + "\n")

                for atom, coords in zip(iso.Atoms, geom):

                    xyz_file.write(atom + " " + str(coords[0]) + " " + str(coords[1]) + " " + str(coords[2]) + "\n")

                xyz_file.close()

                #now need to fragment the molecule and generate these representations

                #build ob mol

                obconversion = OBConversion()
                obconversion.SetInFormat("sdf")
                m = OBMol()
                obconversion.ReadFile(m, iso.InputFile)

                os.mkdir(str(OutputFolder / "dp5" /InputFile.stem) + "_" +str(i).zfill(3) + "_fragments")

                mol_fragments(m,str(OutputFolder / "dp5" /InputFile.stem) + "_" +str(i).zfill(3) + "_fragments")

                conf_rep = []

                for xyz_frag in sorted(os.listdir(  str(OutputFolder / "dp5" /InputFile.stem) + "_" +str(i).zfill(3) + "_fragments")):

                    c = qml.Compound(xyz=str(OutputFolder / "dp5" /InputFile.stem) + "_" +str(i).zfill(3) + "_fragments/" + xyz_frag)

                    c.generate_fchl_representation(max_size=54, cut_distance=c_distance)

                    conf_rep.append(c.representation[0])

                dp5Data.AtomReps[-1].append([])

                for C_l in iso.Clabels:

                    ind = int(C_l.split("C")[1])

                    dp5Data.AtomReps[-1][-1].append(conf_rep[ind])

    return dp5Data


def InternalScaling(dp5Data):
    # perform internal scaling process

    # calculate prediction errors

    if len(dp5Data.Cexp[0]) > 0:

        for Cshifts, Cexp in zip(dp5Data.Cshifts, dp5Data.Cexp):
            dp5Data.Cscaled.append(ScaleNMR(Cshifts, Cexp))

    '''

    if len(dp5Data.Hexp[0]) > 0:

        for Hshifts, Hexp in zip(dp5Data.Hshifts, dp5Data.Hexp):
            dp5Data.Hscaled.append(ScaleNMR(Hshifts, Hexp))

        for Hscaled, Hexp in zip(dp5Data.Hscaled, dp5Data.Hexp):
            dp5Data.Hscalederrors.append([Hscaled[i] - Hexp[i] for i in range(0, len(Hscaled))])
    '''
    return dp5Data


def ScaleNMR(calcShifts, expShifts):

    slope, intercept, r_value, p_value, std_err = stats.linregress(expShifts,
                                                                   calcShifts)
    scaled = [(x - intercept) / slope for x in calcShifts]

    return scaled


def BoltzmannWeight_DP5(Isomers,dp5Data):

    for iso,scaled_probs in zip( Isomers, dp5Data.ScaledAtomProbs):

        B_scaled_probs = [0] * len(scaled_probs[0])

        for population, conf_scaled_p in zip(iso.Populations, scaled_probs ):

            for i in range(len(B_scaled_probs)):

                B_scaled_probs[i] += conf_scaled_p[i] * population

        dp5Data.BScaledAtomProbs.append(B_scaled_probs)

    return dp5Data


def Calculate_DP5(dp5Data):

    for scaled_probs in dp5Data.BScaledAtomProbs:

        DP5scaled = 1 - gmean([1 - p_si for p_si in scaled_probs])

        dp5Data.CScaledprobs.append(DP5scaled)

    return dp5Data


def Rescale_DP5(dp5Data,Settings):

    incorrect_kde = pickle.load(open(Path(Settings.ScriptDir) / "i_w_kde_mean_s_0.025.p" ,"rb"))

    correct_kde = pickle.load(open(Path(Settings.ScriptDir) / "c_w_kde_mean_s_0.025.p" ,"rb"))

    i = 0

    for scaled in dp5Data.BScaledAtomProbs:

        dp5Data.BScaledAtomProbs[i] = [ 1 - float(incorrect_kde.pdf(x) / (incorrect_kde.pdf(x) + correct_kde.pdf(x))) for x in scaled  ]

        i += 1

    dp5Data.DP5scaledprobs = [  1 - float(incorrect_kde.pdf(x) / (incorrect_kde.pdf(x) + correct_kde.pdf(x))) for x in dp5Data.CScaledprobs  ]   # Final DP5S

    return dp5Data


def Pickle_res(dp5Data,Settings):

    data_dic = {"Cshifts": dp5Data.Cshifts,
                "Cexp": dp5Data.Cexp,
                "Clabels": dp5Data.Clabels,
                "Hshifts": dp5Data.Hshifts,
                "Hexp": dp5Data.Hexp,
                "Hlabels": dp5Data.Hlabels,
                "Cscaled": dp5Data.Cscaled,
                "Hscaled": dp5Data.Hscaled,
                "ConfCshifts": dp5Data.ConfCshifts,
                "Compounds": dp5Data.Compounds,
                "AtomReps": dp5Data.AtomReps,

                "ScaledAtomProbs": dp5Data.ScaledAtomProbs,

                "BScaledAtomProbs": dp5Data.BScaledAtomProbs,

                "CScaledprobs": dp5Data.CScaledprobs,

                "DP5scaledprobs": dp5Data.DP5scaledprobs}

    pickle.dump(data_dic , open(Path(Settings.OutputFolder) / "dp5" / "data_dic.p","wb"))

    return dp5Data


def UnPickle_res(dp5Data,Settings):

    data_dic =  pickle.load(open(Path(Settings.OutputFolder) / "dp5" / "data_dic.p","rb"))

    dp5Data.Cshifts = data_dic["Cshifts"]
    dp5Data.Cexp = data_dic["Cexp"]
    dp5Data.Clabels = data_dic["Clabels"]
    dp5Data.Hshifts = data_dic["Hshifts"]
    dp5Data.Hexp = data_dic["Hexp"]
    dp5Data.Hlabels = data_dic["Hlabels"]
    dp5Data.Cscaled = data_dic["Cscaled"]
    dp5Data.Hscaled = data_dic["Hscaled"]
    dp5Data.ConfCshifts = data_dic["ConfCshifts"]
    dp5Data.Compounds = data_dic["Compounds"]
    dp5Data.AtomReps = data_dic["AtomReps"]
    dp5Data.ScaledAtomProbs = data_dic["ScaledAtomProbs"]
    dp5Data.BScaledAtomProbs = data_dic["BScaledAtomProbs"]
    dp5Data.CScaledprobs = data_dic["CScaledprobs"]
    dp5Data.DP5scaledprobs = data_dic["DP5scaledprobs"]

    return dp5Data


def PrintAssignment(dp5Data):
    isomer = 0

    for Clabels, Cshifts, Cexp, Cscaled, atom_p in zip(dp5Data.Clabels, dp5Data.Cshifts, dp5Data.Cexp, dp5Data.Cscaled,dp5Data.BScaledAtomProbs):
        dp5Data.output += ("\n\nAssigned C shifts for isomer " + str(isomer + 1) + ": ")

        PrintNMR(Clabels, Cshifts, Cscaled, Cexp,atom_p, dp5Data)

        isomer += 1


def PrintNMR(labels, values, scaled, exp,atom_p, dp5Data):

    s = np.argsort(values)

    svalues = np.array(values)[s]

    slabels = np.array(labels)[s]

    sscaled = np.array(scaled)[s]

    sexp = np.array(exp)[s]

    atom_p = np.array(atom_p)[s]

    dp5Data.output += ("\nlabel, calc, corrected, exp, error,prob")

    for i in range(len(labels)):

        dp5Data.output += ("\n" + format(slabels[i], "6s") + ' ' + format(svalues[i], "6.2f") + ' '
                           + format(sscaled[i], "6.2f") + ' ' + format(sexp[i], "6.2f") + ' ' +
                           format(sexp[i] - sscaled[i], "6.2f")+ ' ' +
                           format(atom_p[i] , "6.2f"))


def MakeOutput(dp5Data, Isomers, Settings):
    # add some info about the calculation

    dp5Data.output += Settings.InputFiles[0] + "\n"

    dp5Data.output += "\n" + "Solvent = " + Settings.Solvent

    dp5Data.output += "\n" + "Force Field = " + Settings.ForceField + "\n"

    if 'o' in Settings.Workflow:
        dp5Data.output += "\n" + "DFT optimisation Functional = " + Settings.oFunctional
        dp5Data.output += "\n" + "DFT optimisation Basis = " + Settings.oBasisSet

    if 'e' in Settings.Workflow:
        dp5Data.output += "\n" + "DFT energy Functional = " + Settings.eFunctional
        dp5Data.output += "\n" + "DFT energy Basis = " + Settings.eBasisSet

    if 'n' in Settings.Workflow:
        dp5Data.output += "\n" + "DFT NMR Functional = " + Settings.nFunctional
        dp5Data.output += "\n" + "DFT NMR Basis = " + Settings.nBasisSet

    if Settings.StatsParamFile != "none":
        dp5Data.output += "\n\nStats model = " + Settings.StatsParamFile

    dp5Data.output += "\n\nNumber of isomers = " + str(len(Isomers))

    c = 1

    for i in Isomers:
        dp5Data.output += "\nNumber of conformers for isomer " + str(c) + " = " + str(len(i.Conformers))

        c += 1

    PrintAssignment(dp5Data)

    dp5Data.output += ("\n\nResults of DP5 using Carbon: ")

    for i, p in enumerate(dp5Data.DP5scaledprobs):

        dp5Data.output += ("\nIsomer " + str(i + 1) + ": " + format(p * 100, "4.1f") + "%")

    print(dp5Data.output)

    if Settings.OutputFolder == '':

        out = open(str(os.getcwd()) + "/" + str(Settings.InputFiles[0] + "NMR.dp5"), "w+")

    else:

        out = open(os.path.join(Settings.OutputFolder, str(Settings.InputFiles[0] + "NMR.dp5")), "w+")

    out.write(dp5Data.output)

    out.close()

    return dp5Data


def mol_fragments(mole,outfile):

    obconv = OBConversion()

    obconv.SetOutFormat("xyz")

    c = 1

    for atom, exp,dft,diff in zip(OBMolAtomIter(mole)):

        # if this is a carbon atom start a breadth first search for other carbon atoms with depth specified

        # create a new mol instance

        new_mol = OBMol()

        # add this atom

        # new_mol.AddAtom(atom)

        fragment_ind = []

        l = atom.GetIndex()

        fragment_ind.append(l)

        # for iteration depth radius

        old_queue = [atom]

        for iteration in range(0, 3):

            new_queue = []

            for a in old_queue:

                for atom2 in OBAtomAtomIter(a):

                    i = atom2.GetIndex()

                    # if the atom has not been seen before add it to the fragment ind list and to the new molecule

                    if i not in fragment_ind:
                        new_queue.append(atom2)

                        fragment_ind.append(i)

                        # new_mol.AddAtom(atom2)

            old_queue = copy.copy(new_queue)

        fragment_ind = [fragment_ind[0]] + sorted(fragment_ind[1:])

        for i in fragment_ind:

            for a in OBMolAtomIter(mole):

                if a.GetIndex() == i:
                    new_mol.AddAtom(a)

        f = open(outfile + "frag" + str(l).zfill(3) + ".xyz", "w+")

        f.write(str(new_mol.NumAtoms()) + "\n\n")

        i = 0

        for atom in OBMolAtomIter(new_mol):

            f.write(atom.GetType()[0] + " " + str(atom.GetX()) + " " + str(atom.GetY()) + " " + str(
                atom.GetZ()) + "\n")


            i+=1

        f.close()

        c += 1