<jupyter_start><jupyter_text>WBO Interpolation Openforcefield is interested in implementation of wbo interpolated parameters. These torsion parameters calculate the wbo between the two cetral atoms invovled in a torsion. WBO correlates with the torsion barrier height and also captures the behavior of remote chemical properties which allowing for potential reduction of parameters and improved accuracy.  Calculation of k- For each molecule the torsion is checked to match TIG-fit0, and then checks if the application of the torsion is 4-fold around the torsion match- The script sets up a forcebalance optimization for a single torsion targets and performs the optimization- The k value is parsed from the .offxml optimized torsion TIG-fit0<jupyter_code>#imports 
import os
import json
import tempfile
from openforcefield.typing.engines.smirnoff import (ForceField,
UnassignedValenceParameterException, BondHandler, AngleHandler,
ProperTorsionHandler, ImproperTorsionHandler,
vdWHandler)
from plot_td_energies import collect_td_targets_data, plot_td_targets_data
from fragmenter.utils import HARTREE_2_KJMOL
import pickle
import shutil
from openforcefield.topology import Molecule, Topology
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
force_balance_file = 'optimize.in'
from make_torsion_target_new import *
from openforcefield.topology import Molecule, Topology
import pickle
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
import qcportal as ptl
import gzip
def genData(ff, ff2):
    """
    Definition: Generate the kvalue from forcebalance runs for each molecule in the target directory 
    
    ff: The force field used for forcebalance runs, includes "TIG-fit0" parameter 
    ff2: The force field used for labeling torsion parameters for analysis, for example force field with full wbo interpolated parameters 
    """
    #clear out target directory to create new targets based on dataset specifications 
    #shutil.rmtree(target_directory)
    #os.mkdir(target_directory)
    ff_path = '/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/openff-1.3.0.offxml'
    directory = './targets'

    sub_dir = [x[0] for x in os.walk(directory) if os.path.basename(x[0]).startswith('td')]

    forcefield = ForceField(ff_path, allow_cosmetic_attributes=True)        
        
    #Create targets folder in subdirectory 
    force_balance_file = 'optimize.in'


    # Find which line of optimize.in the targets specificaiton begins at
    with open(force_balance_file, 'r') as f:
        replace_index = None
        force_balance_lines = f.readlines()
        for index,line in enumerate(force_balance_lines):
            try:
                #if 'name' in line[:4]:
                if line.startswith('name'):
                    replace_index = index
                    break
            except:
                pass
        #print(replace_index)
    
    count = 0
    plot_dict = {}
    kwbo={}
    kwbo_dict={}
    for asub in sub_dir:
        #do somehting with metadata
        metadata = os.path.join(asub,'metadata.json')
        #do something with mol2 file
        mol2 = os.path.join(asub, 'input.mol2')
        #get indices from json
        with open(metadata, 'r') as f:
            data = json.load(f)
            #print(data['attributes']['canonical_isomeric_explicit_hydrogen_mapped_smiles'])
            cmiles=data['attributes']['canonical_isomeric_explicit_hydrogen_mapped_smiles']
            dihedral=tuple(data['dihedrals'][0])
            entry_value = data['entry_label']
            dataset = data['dataset_name']
            #print(dihedral)
        
        client = ptl.FractalClient()
        ds = client.get_collection("TorsionDriveDataset", dataset)
        ds.status(status="COMPLETE")
        entry = ds.get_entry(entry_value)
        TID = entry.object_map['default']
        print("entry:")
        print(TID)

        #molecules =  Molecule.from_file(mol2, allow_undefined_stereo=True)
        print(cmiles)
        molecules=Molecule.from_mapped_smiles(cmiles)
        topology = Topology.from_molecules([molecules])
        molecules.visualize()
        # Let's label using the Parsley force field
        forcefield = ForceField(ff, allow_cosmetic_attributes=True)

        # Run the molecule labeling
        molecule_force_list = forcefield.label_molecules(topology)
        #print(dict(molecule_force_list[0]['ProperTorsions']))
        # Print out a formatted description of the torsion parameters applied to this molecule
        plot_dict = {}
        for mol_idx, mol_forces in enumerate(molecule_force_list):
            for force_tag, force_dict in mol_forces.items():
                if force_tag == 'ProperTorsions':
                    tigcount=0
                    for (atom_indices, parameter) in force_dict.items():
                        #print('param id', parameter.id)
                        if (parameter.id == "TIG-fit0"):
                            tigcount+=1
                    print(tigcount)
                    #if tigcount==4: 
                    if tigcount>0:
                        for (atom_indices, parameter) in force_dict.items():
                            #print('param id', parameter.id)
                            if (parameter.id == "TIG-fit0"):
                                #print('check1', parameter.id)
                                #check dihedral in forward or reverse order
                                if (atom_indices == dihedral) or (atom_indices == dihedral[::-1]):
                                    #print(atom_indices, dihedral)
                                    #print(asub.split('/')[-1])
                                    #paste
                                    print('check2')
                                    #run forcebalance
                                    tmp_file = tempfile.NamedTemporaryFile(suffix='.in')
                                    with open('debug.in','w') as f1:
                                        for index,line in enumerate(force_balance_lines):
                                            if index != replace_index:
                                                print(line)
                                                f1.write(line)
                                            else:
                                                f1.write(line.split()[0] + ' ' + os.path.basename(asub)+"\n")
                                    print('ForceBalance ' + 'debug.in')
                                    print(os.system('ForceBalance ' + 'debug.in'))
                                    print('done')
                                    
                                    

                                    
                                    #getwbo
                                    wborder= getWBO(cmiles, atom_indices)
                                    print(wborder)
                                    #get kval
                                    #data stored into kval.pickle
                                    kval = readKval(name=os.path.basename(asub)) 
                                    print(kval)
                                    print('kval', kval[os.path.basename(asub)])
                                    #get qdata and store energies
                                    elist = []
                                    with open(os.path.join(asub,'qdata.txt')) as f:
                                        lines = f.readlines()
                                        elist = [float(line.split()[1]) for line in lines if 'ENERGY' in line]
                                    
                                              
                                    
                                    
                                    
                                    
                                    #get the alternative TIG ID 
                                    #added
                                    # Let's label using the Parsley force field
                                    forcefield2 = ForceField(ff2, allow_cosmetic_attributes=True)
                                    # Run the molecule labeling
                                    molecule_force_list = forcefield2.label_molecules(topology)
                                    #print(dict(molecule_force_list[0]['ProperTorsions']))
                                    # Print out a formatted description of the torsion parameters applied to this molecule
                                    #plot_dict = {}
                                    for mol_idx, mol_forces in enumerate(molecule_force_list):
                                        for force_tag, force_dict in mol_forces.items():
                                            if force_tag == 'ProperTorsions':
                                                for (atom_indices, parameter) in force_dict.items():
                                                    if (atom_indices == dihedral) or (atom_indices == dihedral[::-1]):
                                                        parameter_match=parameter.id 
                                                        print(parameter_match)
                                    
                                    
                                    #store info into dictionary for plotting later
                                    kwbo[cmiles]=[wborder,kval[os.path.basename(asub)], parameter_match, getTB(elist), TID]
                                    kwbo_dict[cmiles]={'wbo': wborder, 'kval': kval[os.path.basename(asub)], 'id': parameter_match, 'tb':getTB(elist), 'tid':TID}
                                    #make plot
                                    plot_data = collect_td_targets_data('debug.tmp', 'targets')
                                    plot_dict.update(plot_data)
                                    plot_td_targets_data(plot_data)
    
    
    with open('plot_data_plot_nonring.pk', 'wb') as pk:
        pickle.dump(plot_dict,pk)
    
    with open('plot_k_wbo_list_nonring.pk', 'wb') as pk:
        pickle.dump(kwbo,pk)
        
    with open('plot_k_wbo_dict_nonring.pk', 'wb') as pk:
        pickle.dump(kwbo_dict,pk)

    
    return 


def readKval(read_dir='debug.tmp',modification='TIG-fit0',parameter='k1', name='*'):
    ####Add this to script 
    #from openforcefield.typing.engines.smirnoff import ForceField
    #ff = ForceField('openff-1.0.0.offxml')
    #proper_torsion_handler = ff.get_parameter_handler('ProperTorsions')
    #param = proper_torsion_handler.get_parameter({'id':'t5'})
    #print(param[0].k1)
    import glob
    import pickle
    folders = glob.glob(os.path.join(read_dir,name))
    storage_dict = {}
    for adir in folders:
        print('adir', adir)
        try:
            last_iter = glob.glob((os.path.join(adir,'iter*')))[-1]
            pass_var = 1
        except:
            pass_var = 0
        if pass_var:
            xml = os.path.join(last_iter,'test.offxml')
            with open(xml, 'r') as f:
                lines = f.readlines()
                for line in lines:
                    if modification in line:
                        split = line.split()
                        for asplit in split:
                            if parameter+'=' in asplit:
                                value = float(asplit.split('"')[1])
                                print(os.path.basename(adir))
                                storage_dict[os.path.basename(adir)] = value
    with open('kval.pickle', 'wb') as f:
        pickle.dump(storage_dict, f)
    return storage_dict



def getWBO(cmiles, dihedral_indices):
    """
    Description: Get the WBO of a molecule using the cmiles and dihedral indicies 
    
    input:
    cmiles: String a cmiles for a particular molecule
    dihedral_indices: List of the dihedral indices from a torsion scan 
    
    return:
    wbo: The WBO of a the central bond 

    """
    from openforcefield.topology import Molecule, Topology

    #print(def(Molecule))
    print(dir(Molecule))
    print(cmiles)
    offmol=Molecule.from_mapped_smiles(cmiles)
    #offmol = Molecule.from_mapped_smiles(cmiles)
    offmol.assign_fractional_bond_orders()
    bond = offmol.get_bond_between(dihedral_indices[1], dihedral_indices[2])
    wbo=bond.fractional_bond_order
    return wbo





def getTB(energies):
    """
    Definition: Calculate the torsion barrier from a list of energies from a list of final energies 
    
    input:
    energies: List of final energies
    
    return: 
    torsion barrier: The TB from a torsion scan 
    """
    #energies = list(tdr_object.get_final_energies().values())
    #tmp = list(tdr_object.get_final_energies().keys())
    tmp = energies
    print(tmp)
    start = -180
    angles = [i * np.pi / 180 for i in range(0, len(tmp))]
    print(angles)
    angles, energies = zip(*sorted(zip(angles, energies)))
    angles = np.array(angles)
    energies = np.array(energies)
    angles = np.append(
        angles[-3:] - 2 * np.pi, np.append(angles, angles[:3] + 2 * np.pi)
    )
    energies = np.append(energies[-3:], np.append(energies, energies[:3]))

    idx = []
    for i in range(len(angles) - 2):
        m1 = (energies[i + 1] - energies[i]) / (angles[i + 1] - angles[i])
        m2 = (energies[i + 2] - energies[i + 1]) / (angles[i + 2] - angles[i + 1])
        if np.sign(m1) == np.sign(m2):
            continue
        else:
            idx.append(i + 1)

    torsion_barriers = []
    for i in range(int(len(idx) - 1)):
        torsion_barriers.append(
            abs(HARTREE_2_KJMOL * abs(energies[idx[i]] - energies[idx[i + 1]]))
        )
    torsion_barriers = np.array(torsion_barriers)

    return max(torsion_barriers)




def makeTDTargets(ds, ff):
    forcefield = ForceField(ff, allow_cosmetic_attributes=True)
    #make td Targets 
    for d in ds: 
        torsiondrive_data=download_torsiondrive_data(d)
        #make_torsiondrive_target(d, torsiondrive_data, test_ff=forcefield)
        #with open('torsiondrive_data.pickle', 'rb') as f:
        #    torsiondrive_data = pickle.load(f)
        make_torsiondrive_target(d, torsiondrive_data, test_ff=forcefield)




        
        


#directory = '/Users/jessica/Downloads/release_1.2.0/fb-fit/targets/'
#ff = '/Users/jessica/Documents/Grad_research/WBO_Torsions_ChayaPaper/release_1.3.0_2/fb-fit/fb-fit0/forcefield/test.offxml'

#be very careful when specifying this 
target_directory='/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/targets'
#all of the molecules there is torsions 
#break them out into subplots for the TIG parameters 
#Plot for molecules that use TIG5a, but use the TIG-fit0 plots 
#Plot for TIG3 
#run the same experiments with Carbon-Nitrogen central bonds parameters 
#makePlots(ds, ff, param_fit, ff_analysis, param_analysis, target_directory)


ds=['SMIRNOFF Coverage Torsion Set 1']
param_fit=['TIG-fit0'] 
param_analysis=[]
ff_analysis=''

ff='/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/forcefield/test.offxml'







ff='/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/forcefield/test.offxml'
ff2='/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/forcefield/fit0_optimized.offxml'

genData(ff, ff2)
ds = ['OpenFF Substituted Phenyl Set 1']
#makeTDTargets(ds, ff, param_fit, ff_analysis, param_analysis, target_directory)
#'Fragment Stability Benchmark',
#'OpenFF Fragmenter Validation 1.0',
'Fragment Stability Benchmark',
'OpenFF Fragmenter Validation 1.0',
'OpenFF Gen 2 Torsion Set 1 Roche 2',
'OpenFF Gen 2 Torsion Set 2 Coverage 2',

#problem with this dataset
#'OpenFF Full TorsionDrive Benchmark 1'
'OpenFF Gen 2 Torsion Set 3 Pfizer Discrepancy 2',

ds = [
'OpenFF Gen 2 Torsion Set 4 eMolecules Discrepancy 2',
'OpenFF Gen 2 Torsion Set 5 Bayer 2',
'OpenFF Gen 2 Torsion Set 6 Supplemental 2',
'OpenFF Group1 Torsions',
'OpenFF Group1 Torsions 2',
'OpenFF Group1 Torsions 3',
'OpenFF Primary Benchmark 1 Torsion Set',
'OpenFF Primary Benchmark 2 Torsion Set',
'OpenFF Primary TorsionDrive Benchmark 1',
'OpenFF Rowley Biaryl v1.0',
'OpenFF Substituted Phenyl Set 1',
'OpenFF-benchmark-ligand-fragments-v1.0',
'Pfizer Discrepancy Torsion Dataset 1',
'SMIRNOFF Coverage Torsion Set 1',
'TorsionDrive Paper']

ds = ['OpenFF Substituted Phenyl Set 1']
#ds = ['OpenFF Fragmenter Validation 1.0']
#ds=['OpenFF Full TorsionDrive Benchmark 1']
ff= '/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/forcefield/test.offxml'
#makeTDTargets(ds, ff)

makePlots(ds, ff, param_fit, ff_analysis, param_analysis, target_directory, ff2)
ff2='/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/forcefield/fit0_optimized.offxml'

ff_path = '/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/openff-1.3.0.offxml'
#ff_path = 'openff-1.3.0.offxml'
#ff_path ='/Users/jessica/Documents/Grad_research/wbointerpolation/by_molecule_experiment/fb-fit0/forcefield/test.offxml'
forcefield = ForceField(ff_path, allow_cosmetic_attributes=True)



def plot_wbo_vs_k(filename):
    with open(filename, 'rb') as f:
        x = pickle.load(f)
    #print(x)
    
    wbo=[]
    kval=[]
    ids=set()
    
    for key, item in x.items():
        wbo.append(item[0])
        kval.append(item[1])
        ids.add(item[2])
        
    #print(ids)
    
    dataDictWBO={}
    dataDictKval={}
    
    for i in ids:
        dataDictWBO[i]=[]
        dataDictKval[i]=[] 
        
        
    for i in ids:
        for key, item in x.items():
            if item[2]==i:
                dataDictWBO[i].append(item[0])
                dataDictKval[i].append(item[1])
    
    #print(wbo)
    #print(kval)
    for k, item in dataDictWBO.items():
        plt.scatter(dataDictWBO[k], dataDictKval[k], label=k)
    #plt.scatter(wbo, kval)
    plt.xlabel("WBO")
    plt.legend()
    plt.ylabel("kval (kcal/mol)")
    plt.title("Substituted Phenyl Set: [*:1]~[#6X3:2]~!@[#6X3:3]~[*:4]")
    plt.savefig('initialplot_energycompare.pdf')
    plt.show()
    
def plot_wbo_tb_compare(filename):
    
    with open(filename, 'rb') as f:
        x = pickle.load(f)
    #print(x)
    
    wbo=[]
    kval=[]
    tb=[]
    ids=set()
    
    for key, item in x.items():
        wbo.append(item[0])
        kval.append(item[1])
        ids.add(item[2])
        tb.append(item[3]*0.239006)
    #print(ids)
    print(tb)
    print(kval)
    dataDictWBO={}
    dataDictKval={}
    
    for i in ids:
        dataDictWBO[i]=[]
        dataDictKval[i]=[] 
        
        
    for i in ids:
        for key, item in x.items():
            if item[2]==i:
                dataDictWBO[i].append(item[0])
                dataDictKval[i].append(item[1])

    fig, ax1 = plt.subplots()
    colors = list(matplotlib.colors.cnames.values())
    print(colors)
    colors=['#D3D3D3', '#90EE90', '#D3D3D3', '#FFB6C1', '#FFA07A', '#20B2AA', '#87CEFA', '#778899', '#778899', '#B0C4DE', '#FFFFE0', '#00FF00', '#32CD32', '#FAF0E6', '#FF00FF', '#800000', '#66CDAA', '#0000CD', '#BA55D3', '#9370DB', '#3CB371', '#7B68EE', '#00FA9A', '#48D1CC', '#C71585', '#191970', '#F5FFFA', '#FFE4E1', '#FFE4B5', '#FFDEAD', '#000080', '#FDF5E6', '#808000', '#6B8E23', '#FFA500', '#FF4500', '#DA70D6', '#EEE8AA', '#98FB98', '#AFEEEE', '#DB7093', '#FFEFD5', '#FFDAB9', '#CD853F', '#FFC0CB', '#DDA0DD', '#B0E0E6', '#800080', '#663399', '#FF0000', '#BC8F8F', '#4169E1', '#8B4513', '#FA8072', '#F4A460', '#2E8B57', '#FFF5EE', '#A0522D', '#C0C0C0', '#87CEEB', '#6A5ACD', '#708090', '#708090', '#FFFAFA', '#00FF7F', '#4682B4', '#D2B48C', '#008080', '#D8BFD8', '#FF6347', '#40E0D0', '#EE82EE', '#F5DEB3', '#FFFFFF', '#F5F5F5', '#FFFF00', '#9ACD32']
    
    j=3
    for i, val in enumerate(wbo):
        ax1.scatter(wbo[i], kval[i], color=colors[i], marker='o', s=80, alpha=.5)
        
    color = 'tab:red'
    ax1.set_xlabel("WBO")
    ax1.set_ylabel("k (kcal/mol)")
    #ax1.scatter(wbo, kval, color=color)
    ax1.tick_params(axis='y')
    #ax1.legend()
    
    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

    color = 'tab:blue'
    
    #ax2.set_ylabel('sin', color=color)  # we already handled the x-label with ax1
    
    for i, val in enumerate(wbo):
        ax2.scatter(wbo[i], tb[i], color=colors[i], marker='^', s=80, alpha=.5)
    
    #ax2.scatter(wbo, tb, color=color)
    ax2.tick_params(axis='y')
    ax2.set_ylabel('Torsion Barrier (kcal/mol)')
    
    #ax2.legend()
    
    fig.tight_layout()  # otherwise the right y-label is slightly clipped
    
    
    marks = [
                   Line2D([0], [0], marker='o', color='w', label='k',
                          markerfacecolor='grey', markersize=12, alpha=.5),
             Line2D([0], [0], marker='^', color='w', label='TB',
                          markerfacecolor='grey', markersize=12, alpha=.5)]

    plt.legend(handles=marks)
        
    plt.title("Torsion Barrier vs k: [*:1]~[#6X3:2]~!@[#6X3:3]~[*:4] \n Substituted Phenyl Dataset")
    plt.savefig('tbvsk.pdf', bbox_inches='tight')
    
    plt.show()
    
def plot_wbo_tb_compare_solid(filename):
    
    with open(filename, 'rb') as f:
        x = pickle.load(f)
    #print(x)
    
    wbo=[]
    kval=[]
    tb=[]
    ids=set()
    
    for key, item in x.items():
        wbo.append(item[0])
        kval.append(item[1])
        ids.add(item[2])
        tb.append(item[3]*0.239006)
    #print(ids)
    print(tb)
    print(kval)
    dataDictWBO={}
    dataDictKval={}
    
    for i in ids:
        dataDictWBO[i]=[]
        dataDictKval[i]=[] 
        
        
    for i in ids:
        for key, item in x.items():
            if item[2]==i:
                dataDictWBO[i].append(item[0])
                dataDictKval[i].append(item[1])

    fig, ax1 = plt.subplots()
    colors = list(matplotlib.colors.cnames.values())
    
    #for i, val in enumerate(wbo):
    #    ax1.scatter(wbo[i], kval[i], color=colors[i], marker='o', s=50, alpha=.5)
    color = 'tab:red'
    ax1.scatter(wbo, kval, color=color, marker='o', s=80, alpha=.5, label="k")
    ax1.set_xlabel("WBO")
    ax1.set_ylabel("k (kcal/mol)", color=color)
    #ax1.scatter(wbo, kval, color=color)
    ax1.tick_params(axis='y', labelcolor=color)

    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

    color = 'tab:blue'
    #color=['#ffb3ba']
    #ax2.set_ylabel('sin', color=color)  # we already handled the x-label with ax1
    
    
    #ax1.scatter(wbo, tb, color=color, marker='^', s=50, alpha=.5)
    
    ax2.scatter(wbo, tb, color=color, marker='^', s=80, alpha=.5, label="TB")
    ax2.tick_params(axis='y', labelcolor=color)
    ax2.set_ylabel('Torsion Barrier (kcal/mol)', color=color)
    fig.tight_layout()  # otherwise the right y-label is slightly clipped
    elements=[Line2D([0], [0], color='red', label='k',markerfacecolor='red', markersize=15, alpha=.5),Line2D([0], [0], color=color, label='TB', markersize=15, alpha=.5) ]
    plt.legend(handles=elements)
    plt.title("Torsion Barrier vs k: [*:1]~[#6X3:2]~!@[#6X3:3]~[*:4] \n Substituted Phenyl Dataset")
    
    plt.savefig('tbvsk_solidcolor.pdf', bbox_inches='tight')

    plt.show()
    
    

plot_wbo_vs_k('plot_k_wbo.pk')
plot_wbo_tb_compare('plot_k_wbo.pk')
plot_wbo_tb_compare_solid('plot_k_wbo.pk')
        
with gzip.open('sterics_data.pkl', 'rb') as f:
    lj = pickle.load(f)
print(lj.columns)
print(lj['LJ energy diff kcal/mol'])
print('keys',lj['LJ energy diff kcal/mol'].keys())
print('energy',lj['LJ energy diff kcal/mol']['1762227'])


ID=['18537025', '3745392', '3745370', '18886248', '18536199', '18045331', '18535845', '18886245', '18535848', '18537010', '4269710', '18537009', '18537022', '4269708', '4269709', '19953582', '18045328', '18537012', '18537008', '18537018', '4269707', '3745476', '18536057', '18536192', '18045333', '4269712', '18535844', '18886247', '3745617', '18045330', '18537019', '6098536', '18886246', '18886237', '4269703', '18536064', '18535839', '19953583', '6098519', '18886249', '18536061', '19953581', '3745367', '18536065', '4269711', '3745366', '18537017', '4269706', '3745658', '3745390', '18535843', '18537023', '18535849', '4269704', '3745541', '18535838', '19953579', '3745391', '19953580', '18537021', '18886244', '18886238', '18537027', '18537026', '18535846', '18535842', '18886243', '18537020', '18536062', '18536201', '4269705', '18535847', '18536063', '18536056']
#for i in ID:
#    print(lj['LJ energy diff kcal/mol'][i])
pk_keys = lj['LJ energy diff kcal/mol'].keys()
append_list = []
nomatch=[]
for i in ID:
    if i in pk_keys:
        append_list.append(i)
    else:
        nomatch.append(i)
print('match', append_list)
print('no match', nomatch)

        

#print(lj['LJ energy diff kcal/mol']['18537025'])

#kwbo_dict[cmiles]={'wbo': wborder, 'kval': kval[os.path.basename(asub)], 'id': parameter_match, 'tb':getTB(elist), 'tid':TID}


def plot_wbo_tb_compare_solid(filename):
    
    with open(filename, 'rb') as f:
        x = pickle.load(f)
    #print(x)
    
    wbo=[]
    kval=[]
    tb=[]
    ids=set()
    tid=[]
    #print('debug', [(key,item) for key,item in x.items()])
    #print(x)
    tids=['3745392', '3745370', '18535845', '18535848', '4269710', '4269708', '4269709', '19953582', '18045328', '4269707', '3745476', '18536057', '18045333', '4269712', '18886247', '3745617', '18045330', '18537019', '18886237', '4269703', '18536064', '18535839', '19953583', '18886249', '18536061', '19953581', '3745367', '4269711', '3745366', '4269706', '3745658', '3745390', '18535849', '4269704', '3745541', '18535838', '19953579', '3745391', '19953580', '18886244', '18886238', '18535846', '18535842', '18886243', '18536062', '18536201', '4269705', '18535847', '18536063', '18536056']
    
    print(x)
    
    for key, item in x.items():
        if item['tid'] in tids:
            if lj['LJ energy diff kcal/mol'][item['tid']]<3.4:
                #print("Not filtered")
                wbo.append(item['wbo'])
                kval.append(item['kval'])
                ids.add(item['id'])
                tid.append(item['tid'])
                tb.append(item['tb']*0.239006)
    #print(ids)
    #print(tid)
    #print(tb)
    #print(kval)
    dataDictWBO={}
    dataDictKval={}
    
    #print(ids)

    
    for i in ids:
        dataDictWBO[i]=[]
        dataDictKval[i]=[] 
        
        
    #print(x)
    for i in ids:
        for key, item in x.items():
            #print(item)
            #print(key)
            #print(item['id'])
            #print(i)
            if item['id']==i:
                dataDictWBO[i].append(item['wbo'])
                dataDictKval[i].append(item['kval'])

    fig, ax1 = plt.subplots()
    colors = list(matplotlib.colors.cnames.values())
    
    #for i, val in enumerate(wbo):
    #    ax1.scatter(wbo[i], kval[i], color=colors[i], marker='o', s=50, alpha=.5)
    color = 'tab:red'
    ax1.scatter(wbo, kval, color=color, marker='o', s=80, alpha=.5, label="k")
    
    error=[]
    for k, t in zip(kval, tb):
        error.append(k-(t/4))
    print(error)
    ax1.errorbar(wbo, kval, yerr=error, fmt='none', color=color, alpha=.5)
    
    ax1.set_xlabel("WBO")
    ax1.set_ylabel("k (kcal/mol)", color=color)
    #ax1.scatter(wbo, kval, color=color)
    ax1.tick_params(axis='y', labelcolor=color)

    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

    color = 'tab:blue'
    #color=['#ffb3ba']
    #ax2.set_ylabel('sin', color=color)  # we already handled the x-label with ax1
    
    
    #ax1.scatter(wbo, tb, color=color, marker='^', s=50, alpha=.5)
    
    ax2.scatter(wbo, tb, color=color, marker='^', s=80, alpha=.5, label="TB")
    ax2.tick_params(axis='y', labelcolor=color)
    ax2.set_ylabel('Torsion Barrier (kcal/mol)', color=color)
    fig.tight_layout()  # otherwise the right y-label is slightly clipped
    elements=[Line2D([0], [0], color='red', label='k',markerfacecolor='red', markersize=15, alpha=.5),Line2D([0], [0], color=color, label='TB', markersize=15, alpha=.5) ]
    
    
    
    plt.legend(handles=elements)
    plt.title("Torsion Barrier vs k: [*:1]~[#6X3:2]~!@[#6X3:3]~[*:4] \n 1.2.0 training & subtituted phenyl datasets")
    
    plt.savefig('tbvsk_solidcolor.pdf', bbox_inches='tight')

    plt.show()
    
def plot_wbo_tb_compare_solid_noerrorbar(filename):
    with open(filename, 'rb') as f:
        x = pickle.load(f)
    #print(x)
    
    wbo=[]
    kval=[]
    tb=[]
    ids=set()
    tid=[]
    #print('debug', [(key,item) for key,item in x.items()])
    #print(x)
    tids=['3745392', '3745370', '18535845', '18535848', '4269710', '4269708', '4269709', '19953582', '18045328', '4269707', '3745476', '18536057', '18045333', '4269712', '18886247', '3745617', '18045330', '18537019', '18886237', '4269703', '18536064', '18535839', '19953583', '18886249', '18536061', '19953581', '3745367', '4269711', '3745366', '4269706', '3745658', '3745390', '18535849', '4269704', '3745541', '18535838', '19953579', '3745391', '19953580', '18886244', '18886238', '18535846', '18535842', '18886243', '18536062', '18536201', '4269705', '18535847', '18536063', '18536056']

    for key, item in x.items():
        if item['tid'] in tids:
            if lj['LJ energy diff kcal/mol'][item['tid']]<3.4:
                #print("Not filtered")
                wbo.append(item['wbo'])
                kval.append(item['kval'])
                ids.add(item['id'])
                tid.append(item['tid'])
                tb.append(item['tb']*0.239006)
    #print(ids)
    #print(tid)
    #print(tb)
    #print(kval)
    dataDictWBO={}
    dataDictKval={}
    
    #print(ids)

    
    for i in ids:
        dataDictWBO[i]=[]
        dataDictKval[i]=[] 
        
        
    #print(x)
    for i in ids:
        for key, item in x.items():
            #print(item)
            #print(key)
            #print(item['id'])
            #print(i)
            if item['id']==i:
                dataDictWBO[i].append(item['wbo'])
                dataDictKval[i].append(item['kval'])

    fig, ax1 = plt.subplots()
    colors = list(matplotlib.colors.cnames.values())
    
    #for i, val in enumerate(wbo):
    #    ax1.scatter(wbo[i], kval[i], color=colors[i], marker='o', s=50, alpha=.5)
    color = 'tab:red'
    ax1.scatter(wbo, kval, color=color, marker='o', s=80, alpha=.5, label="k")
    
    #ax1.errorbar(wbo, kval, yerr=error, fmt='none', color=color, alpha=.5)
    
    ax1.set_xlabel("WBO")
    ax1.set_ylabel("k (kcal/mol)", color=color)
    #ax1.scatter(wbo, kval, color=color)
    ax1.tick_params(axis='y', labelcolor=color)

    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

    color = 'tab:blue'
    #color=['#ffb3ba']
    #ax2.set_ylabel('sin', color=color)  # we already handled the x-label with ax1
    
    
    #ax1.scatter(wbo, tb, color=color, marker='^', s=50, alpha=.5)
    
    ax2.scatter(wbo, tb, color=color, marker='^', s=80, alpha=.5, label="TB")
    ax2.tick_params(axis='y', labelcolor=color)
    ax2.set_ylabel('Torsion Barrier (kcal/mol)', color=color)
    fig.tight_layout()  # otherwise the right y-label is slightly clipped
    elements=[Line2D([0], [0], color='red', label='k',markerfacecolor='red', markersize=15, alpha=.5),Line2D([0], [0], color=color, label='TB', markersize=15, alpha=.5) ]
    
    
    
    plt.legend(handles=elements)
    plt.title("Torsion Barrier vs k: [*:1]~[#6X3:2]~!@[#6X3:3]~[*:4] \n 1.2.0 training & subtituted phenyl datasets")
    
    plt.savefig('tbvsk_solidcolor_noerror_bar_steric_filter.pdf', bbox_inches='tight')

    plt.show()  

plot_wbo_tb_compare_solid('plot_k_wbo_dict.pk')
plot_wbo_tb_compare_solid_noerrorbar('plot_k_wbo_dict.pk')
with gzip.open('sterics_data.pkl', 'rb') as f:
    lj = pickle.load(f)
print(lj.columns)
print(lj['LJ energy diff kcal/mol'])
print('keys',lj['LJ energy diff kcal/mol'].keys())
print('energy',lj['LJ energy diff kcal/mol']['1762227'])


ID=['18537025', '3745392', '3745370', '18886248', '18536199', '18045331', '18535845', '18886245', '18535848', '18537010', '4269710', '18537009', '18537022', '4269708', '4269709', '19953582', '18045328', '18537012', '18537008', '18537018', '4269707', '3745476', '18536057', '18536192', '18045333', '4269712', '18535844', '18886247', '3745617', '18045330', '18537019', '6098536', '18886246', '18886237', '4269703', '18536064', '18535839', '19953583', '6098519', '18886249', '18536061', '19953581', '3745367', '18536065', '4269711', '3745366', '18537017', '4269706', '3745658', '3745390', '18535843', '18537023', '18535849', '4269704', '3745541', '18535838', '19953579', '3745391', '19953580', '18537021', '18886244', '18886238', '18537027', '18537026', '18535846', '18535842', '18886243', '18537020', '18536062', '18536201', '4269705', '18535847', '18536063', '18536056']
#for i in ID:
#    print(lj['LJ energy diff kcal/mol'][i])
pk_keys = lj['LJ energy diff kcal/mol'].keys()
append_list = []
nomatch=[]
for i in ID:
    if i in pk_keys:
        append_list.append(i)
    else:
        nomatch.append(i)
print('match', append_list)
print('no match', nomatch)

        

#print(lj['LJ energy diff kcal/mol']['18537025'])

#kwbo_dict[cmiles]={'wbo': wborder, 'kval': kval[os.path.basename(asub)], 'id': parameter_match, 'tb':getTB(elist), 'tid':TID}


def plot_wbo_tb_compare_solid(filename):
    
    with open(filename, 'rb') as f:
        x = pickle.load(f)
    #print(x)
    
    wbo=[]
    kval=[]
    tb=[]
    ids=set()
    tid=[]
    #print('debug', [(key,item) for key,item in x.items()])
    #print(x)
    tids=['3745392', '3745370', '18535845', '18535848', '4269710', '4269708', '4269709', '19953582', '18045328', '4269707', '3745476', '18536057', '18045333', '4269712', '18886247', '3745617', '18045330', '18537019', '18886237', '4269703', '18536064', '18535839', '19953583', '18886249', '18536061', '19953581', '3745367', '4269711', '3745366', '4269706', '3745658', '3745390', '18535849', '4269704', '3745541', '18535838', '19953579', '3745391', '19953580', '18886244', '18886238', '18535846', '18535842', '18886243', '18536062', '18536201', '4269705', '18535847', '18536063', '18536056']
    
    #print(x)
    
    
    
    for key, item in x.items():
        if item['tid'] in tids:
            if lj['LJ energy diff kcal/mol'][item['tid']]<3.4:
                #print("Not filtered")
                wbo.append(item['wbo'])
                kval.append(item['kval'])
                ids.add(item['id'])
                tid.append(item['tid'])
                tb.append(item['tb']*0.239006)
    #print(ids)
    #print(tid)
    #print(tb)
    #print(kval)
    dataDictWBO={}
    dataDictKval={}
    
    #print(ids)

    
    for i in ids:
        dataDictWBO[i]=[]
        dataDictKval[i]=[] 
        
        
    #print(x)
    for i in ids:
        for key, item in x.items():
            #print(item)
            #print(key)
            #print(item['id'])
            #print(i)
            if item['id']==i:
                dataDictWBO[i].append(item['wbo'])
                dataDictKval[i].append(item['kval'])

    fig, ax1 = plt.subplots()
    colors = list(matplotlib.colors.cnames.values())
    
    #for i, val in enumerate(wbo):
    #    ax1.scatter(wbo[i], kval[i], color=colors[i], marker='o', s=50, alpha=.5)
    color = 'tab:red'
    ax1.scatter(wbo, kval, color=color, marker='o', s=80, alpha=.5, label="k")
    
    error=[]
    for k, t in zip(kval, tb):
        error.append(k-(t/4))
    print(error)
    ax1.errorbar(wbo, kval, yerr=error, fmt='none', color=color, alpha=.5)
    
    ax1.set_xlabel("WBO")
    ax1.set_ylabel("k (kcal/mol)", color=color)
    #ax1.scatter(wbo, kval, color=color)
    ax1.tick_params(axis='y', labelcolor=color)

    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

    color = 'tab:blue'
    #color=['#ffb3ba']
    #ax2.set_ylabel('sin', color=color)  # we already handled the x-label with ax1
    
    
    #ax1.scatter(wbo, tb, color=color, marker='^', s=50, alpha=.5)
    
    ax2.scatter(wbo, tb, color=color, marker='^', s=80, alpha=.5, label="TB")
    ax2.tick_params(axis='y', labelcolor=color)
    ax2.set_ylabel('Torsion Barrier (kcal/mol)', color=color)
    fig.tight_layout()  # otherwise the right y-label is slightly clipped
    elements=[Line2D([0], [0], color='red', label='k',markerfacecolor='red', markersize=15, alpha=.5),Line2D([0], [0], color=color, label='TB', markersize=15, alpha=.5) ]
    
    
    
    plt.legend(handles=elements)
    plt.title("Torsion Barrier vs k: [*:1]~[#6X3:2]~!@[#6X3:3]~[*:4] \n 1.2.0 training & subtituted phenyl datasets")
    
    plt.savefig('tbvsk_solidcolor.pdf', bbox_inches='tight')

    plt.show()
    
def plot_wbo_tb_compare_solid_noerrorbar(filename):
    with open(filename, 'rb') as f:
        x = pickle.load(f)
    #print(x)
    
    wbo=[]
    kval=[]
    tb=[]
    ids=set()
    tid=[]
    #print('debug', [(key,item) for key,item in x.items()])
    #print(x)
    tids=['3745392', '3745370', '18535845', '18535848', '4269710', '4269708', '4269709', '19953582', '18045328', '4269707', '3745476', '18536057', '18045333', '4269712', '18886247', '3745617', '18045330', '18537019', '18886237', '4269703', '18536064', '18535839', '19953583', '18886249', '18536061', '19953581', '3745367', '4269711', '3745366', '4269706', '3745658', '3745390', '18535849', '4269704', '3745541', '18535838', '19953579', '3745391', '19953580', '18886244', '18886238', '18535846', '18535842', '18886243', '18536062', '18536201', '4269705', '18535847', '18536063', '18536056']

    for key, item in x.items():
        if item['tid'] in tids:
            if lj['LJ energy diff kcal/mol'][item['tid']]<3.4:
                #print("Not filtered")
                wbo.append(item['wbo'])
                kval.append(item['kval'])
                ids.add(item['id'])
                tid.append(item['tid'])
                tb.append(item['tb']*0.239006)
    #print(ids)
    #print(tid)
    #print(tb)
    #print(kval)
    dataDictWBO={}
    dataDictKval={}
    
    #print(ids)

    
    for i in ids:
        dataDictWBO[i]=[]
        dataDictKval[i]=[] 
        
        
    #print(x)
    for i in ids:
        for key, item in x.items():
            #print(item)
            #print(key)
            #print(item['id'])
            #print(i)
            if item['id']==i:
                dataDictWBO[i].append(item['wbo'])
                dataDictKval[i].append(item['kval'])

    fig, ax1 = plt.subplots()
    colors = list(matplotlib.colors.cnames.values())
    
    #for i, val in enumerate(wbo):
    #    ax1.scatter(wbo[i], kval[i], color=colors[i], marker='o', s=50, alpha=.5)
    color = 'tab:red'
    ax1.scatter(wbo, kval, color=color, marker='o', s=80, alpha=.5, label="k")
    
    #ax1.errorbar(wbo, kval, yerr=error, fmt='none', color=color, alpha=.5)
    
    ax1.set_xlabel("WBO")
    ax1.set_ylabel("k (kcal/mol)", color=color)
    #ax1.scatter(wbo, kval, color=color)
    ax1.tick_params(axis='y', labelcolor=color)

    ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis

    color = 'tab:blue'
    #color=['#ffb3ba']
    #ax2.set_ylabel('sin', color=color)  # we already handled the x-label with ax1
    
    
    #ax1.scatter(wbo, tb, color=color, marker='^', s=50, alpha=.5)
    
    ax2.scatter(wbo, tb, color=color, marker='^', s=80, alpha=.5, label="TB")
    ax2.tick_params(axis='y', labelcolor=color)
    ax2.set_ylabel('Torsion Barrier (kcal/mol)', color=color)
    fig.tight_layout()  # otherwise the right y-label is slightly clipped
    elements=[Line2D([0], [0], color='red', label='k',markerfacecolor='red', markersize=15, alpha=.5),Line2D([0], [0], color=color, label='TB', markersize=15, alpha=.5) ]
    
    
    
    plt.legend(handles=elements)
    plt.title("Torsion Barrier vs k: [*:1]~[#6X3:2]~!@[#6X3:3]~[*:4] \n 1.2.0 training & subtituted phenyl datasets")
    
    plt.savefig('tbvsk_solidcolor_noerror_bar_steric_filter.pdf', bbox_inches='tight')

    plt.show()  

plot_wbo_tb_compare_solid('plot_k_wbo_dict.pk')
plot_wbo_tb_compare_solid_noerrorbar('plot_k_wbo_dict.pk')
import gzip
from openeye import oechem

# Highlight element of interest
#subs = oechem.OESubSearch("[#6]") # carbon
subs = None

# Read the input SMILES
with open('optimization_inputs.smi', 'r') as infile:
    smiles_list = infile.readlines()
smiles_list = [ smiles.strip() for smiles in smiles_list ]

with open(filename, 'rb') as f:
    x = pickle.load(f)
print(x)


# Build OEMols
oemols = list()
for smiles in smiles_list:
    oemol = oechem.OEMol()
    oechem.OESmilesToMol(oemol, smiles)
    oemols.append(oemol)

# Generate a PDF of all molecules in the set
pdf_filename = 'optimization_inputs.pdf'

from openeye import oedepict
itf = oechem.OEInterface()
PageByPage = True
suppress_h = True
rows = 10
cols = 6
ropts = oedepict.OEReportOptions(rows, cols)
ropts.SetHeaderHeight(25)
ropts.SetFooterHeight(25)
ropts.SetCellGap(2)
ropts.SetPageMargins(10)
report = oedepict.OEReport(ropts)
cellwidth, cellheight = report.GetCellWidth(), report.GetCellHeight()
opts = oedepict.OE2DMolDisplayOptions(cellwidth, cellheight, oedepict.OEScale_Default * 0.5)
opts.SetAromaticStyle(oedepict.OEAromaticStyle_Circle)
pen = oedepict.OEPen(oechem.OEBlack, oechem.OEBlack, oedepict.OEFill_On, 1.0)
opts.SetDefaultBondPen(pen)
oedepict.OESetup2DMolDisplayOptions(opts, itf)
for i, mol in enumerate(oemols):
    cell = report.NewCell()
    mol_copy = oechem.OEMol(mol)
    oedepict.OEPrepareDepiction(mol_copy, False, suppress_h)
    disp = oedepict.OE2DMolDisplay(mol_copy, opts)

    # Highlight element of interest
    unique = False
    if subs is not None:
        for match in subs.Match(mol_copy, unique):
            oedepict.OEAddHighlighting(disp, oechem.OEColor(oechem.OEYellow), oedepict.OEHighlightStyle_BallAndStick, match)

    oedepict.OERenderMolecule(cell, disp)
    #oedepict.OEDrawCurvedBorder(cell, oedepict.OELightGreyPen, 10.0)

oedepict.OEWriteReport(pdf_filename, report)


heavyAtomCount=[]
atoms=[]
for mol in oemols:
    heavyAtomCount.append(oechem.OECount(mol, oechem.OEIsHeavy()))
#    atoms.append(oechem.OECount(mol, oechem.GetExplicitDegree()))
print(heavyAtomCount)
print(min(heavyAtomCount))
print(max(heavyAtomCount))


#print(atoms)
#print(min(atoms))
#print(max(atoms))
adict = readKval()
print(adict)
import xml
print(dir(xml.etree))
import xml.etree.ElementTree as ET
root = ET.parse('debug.tmp/td_OpenFF_Substituted_Phenyl_Set_1_152_C11ClH8NO/iter_0001/test.offxml').getroot()
all_name_elements = tree.findall('*/name')

tree = ET.parse('debug.tmp/td_OpenFF_Substituted_Phenyl_Set_1_152_C11ClH8NO/iter_0001/test.offxml').getroot()
all_name_elements = tree.findall('*/TIG-fit0')
print(all_name_elements)
print(dir(tree))
cd ../../
plot_data = collect_td_targets_data('debug.tmp', 'targets')<jupyter_output>match 25
target 25
Collecting td data from 25 target folders
0 td_OpenFF_Substituted_Phenyl_Set_1_009_C11H13NO3
1 td_OpenFF_Substituted_Phenyl_Set_1_010_C11H17N2O2
2 td_OpenFF_Substituted_Phenyl_Set_1_013_C10H13NO3
3 td_OpenFF_Substituted_Phenyl_Set_1_032_C8H10N2O2
4 td_OpenFF_Substituted_Phenyl_Set_1_033_C9H10O3
5 td_OpenFF_Substituted_Phenyl_Set_1_034_C9H12N2O2
6 td_OpenFF_Substituted_Phenyl_Set_1_056_C7H8N2O2
7 td_OpenFF_Substituted_Phenyl_Set_1_084_C9H13N2O2
8 td_OpenFF_Substituted_Phenyl_Set_1_131_C7H4N2O2
9 td_OpenFF_Substituted_Phenyl_Set_1_139_C8H8N2O3
10 td_OpenFF_Substituted_Phenyl_Set_1_141_C11H10NO
11 td_OpenFF_Substituted_Phenyl_Set_1_142_C11H9NO
12 td_OpenFF_Substituted_Phenyl_Set_1_143_C11H9NO
13 td_OpenFF_Substituted_Phenyl_Set_1_144_C11H8NO
14 td_OpenFF_Substituted_Phenyl_Set_1_145_C11FH9NO
15 td_OpenFF_Substituted_Phenyl_Set_1_146_C11FH8NO
16 td_OpenFF_Substituted_Phenyl_Set_1_147_C11FH8NO
17 td_OpenFF_Substituted_Phenyl_Set_1_148_C11FH7NO
18 td_OpenFF_Substituted_Phe[...]
{'fraction_non_alphanumeric': 0.14627218397073194, 'fraction_numerical': 0.08405480946213102, 'mean_word_length': 2.716891891891892, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 20}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 86, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '507755', 'license': 'MIT', 'n_tokens_mistral': 18143, 'n_tokens_neox': 15591, 'n_words': 3100, 'path': 'by_molecule_fits/kversusWBOplots_cleanup.ipynb', 'repo_name': 'danieljcabrera/wbointerpolation', 'size': '44007'}