application_vnd.jupyter.stderr_output_1.png
import numpy as np # linear algebra

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

train = pd.read_csv('../input/train.csv')
test = pd.read_csv('../input/test.csv')
sub = pd.read_csv('../input/sample_submission.csv')

structures = pd.read_csv('input/structures.csv')

atomic_radius = {'H': 0.38, 'C': 0.77, 'N': 0.75, 'O': 0.73, 'F': 0.71, np.nan: 0}
atomic_number = {'H': 1, 'C': 6, 'N': 7, 'O': 8, 'F': 9, np.nan: 0}
atomic_mass = {'H': 1.0079, 'C': 12.0107, 'N': 14.0067, 'O': 15.9994, 'F': 18.9984, np.nan: 0}
vanderwaalsradius = {'H': 120, 'C': 185, 'N': 154, 'O': 140, 'F': 135, np.nan: 0}
covalenzradius = {'H': 30, 'C': 77, 'N': 70, 'O': 66, 'F': 58, np.nan: 0}
electronegativity = {'H': 2.2, 'C': 2.55, 'N': 3.04, 'O': 3.44, 'F': 3.98, np.nan: 0}
ionization_energy = {'H': 13.5984, 'C': 11.2603, 'N': 14.5341, 'O': 13.6181, 'F': 17.4228, np.nan: np.inf}

def atom_props(df, suffix):
    df['atomic_radius' + suffix] = df['atom_' + suffix].apply(lambda x: atomic_radius[x])
    df['atomic_protons' + suffix] = df['atom_' + suffix].apply(lambda x: atomic_number[x])
    df['atomic_mass' + suffix] = df['atom_' + suffix].apply(lambda x: atomic_mass[x])
    df['vanderwaalsradius' + suffix] = df['atom_' + suffix].apply(lambda x: vanderwaalsradius[x])
    df['covalenzradius' + suffix] = df['atom_' + suffix].apply(lambda x: covalenzradius[x])
    df['electronegativity' + suffix] = df['atom_' + suffix].apply(lambda x: electronegativity[x])
    df['ionization_energy' + suffix] = df['atom_' + suffix].apply(lambda x: ionization_energy[x])
    return df

train = pd.merge(train, structures[['molecule_name', 'atom_index', 'atom']], how='left', left_on=['molecule_name', 'atom_index_0'], right_on=['molecule_name', 'atom_index']).rename({'atom': 'atom_0'}, axis=1)
train = pd.merge(train, structures[['molecule_name', 'atom_index', 'atom']], how='left', left_on=['molecule_name', 'atom_index_1'], right_on=['molecule_name', 'atom_index']).rename({'atom': 'atom_1'}, axis=1)
test = pd.merge(test, structures[['molecule_name', 'atom_index', 'atom']], how='left', left_on=['molecule_name', 'atom_index_0'], right_on=['molecule_name', 'atom_index']).rename({'atom': 'atom_0'}, axis=1)
test = pd.merge(test, structures[['molecule_name', 'atom_index', 'atom']], how='left', left_on=['molecule_name', 'atom_index_1'], right_on=['molecule_name', 'atom_index']).rename({'atom': 'atom_1'}, axis=1)

atom_cnt = structures['molecule_name'].value_counts().reset_index(level=0)
atom_cnt.rename({'index': 'molecule_name', 'molecule_name': 'atom_count'}, axis=1, inplace=True)
train = pd.merge(train, atom_cnt, how='left', on='molecule_name')
test = pd.merge(test, atom_cnt, how='left', on='molecule_name')
del atom_cnt

def lr(df):
    df['atom_index_0l'] = df['atom_index_0'].apply(lambda i: max(i - 1, 0))
    tmp = df[['atom_index_0', 'atom_count']]
    df['atom_index_0r'] = tmp.apply(lambda row: min(row['atom_index_0'] + 1, row['atom_count']), axis=1)
    df = pd.merge(df, structures[['molecule_name', 'atom_index', 'atom']], how='left', left_on=['molecule_name', 'atom_index_0l'], right_on=['molecule_name', 'atom_index']).rename({'atom': 'atom_0l'}, axis=1)
    df = pd.merge(df, structures[['molecule_name', 'atom_index', 'atom']], how='left', left_on=['molecule_name', 'atom_index_0r'], right_on=['molecule_name', 'atom_index']).rename({'atom': 'atom_0r'}, axis=1)
    df['atom_index_1l'] = df['atom_index_1'].apply(lambda i: max(i - 1, 0))
    tmp = df[['atom_index_1', 'atom_count']]
    df['atom_index_1r'] = tmp.apply(lambda row: min(row['atom_index_1'] + 1, row['atom_count']), axis=1)
    df = pd.merge(df, structures[['molecule_name', 'atom_index', 'atom']], how='left', left_on=['molecule_name', 'atom_index_1l'], right_on=['molecule_name', 'atom_index']).rename({'atom': 'atom_1l'}, axis=1)
    df = pd.merge(df, structures[['molecule_name', 'atom_index', 'atom']], how='left', left_on=['molecule_name', 'atom_index_1r'], right_on=['molecule_name', 'atom_index']).rename({'atom': 'atom_1r'}, axis=1)
    return df

train = lr(train)
test = lr(test)

train = atom_props(train, '0')
train = atom_props(train, '0l')
train = atom_props(train, '0r')
train = atom_props(train, '1')
train = atom_props(train, '1l')
train = atom_props(train, '1r')
test = atom_props(test, '0')
test = atom_props(test, '0l')
test = atom_props(test, '0r')
test = atom_props(test, '1')
test = atom_props(test, '1l')
test = atom_props(test, '1r')