Thursday 24 September 2020

Pandas: interpolate between column values to return index name

I'd like to return the index name and cost in the following code. I think I'm on the right track, but not sure how to push this forward, I've attached my code and output

EDIT: I'm open to design changes in the way I'm doing this.

for example a in test_data, it should return 'a': 'rep_asp_art', '1000000'

import pandas as pd


treatments = {'min_asp_art': {'treat': 'min', 'mat': 'asp', 'class': 'art',
                              'hci': 85, 'cost': 100000},
              'maj_asp_art': {'treat': 'maj', 'mat': 'asp', 'class': 'art',
                              'hci': 60, 'cost': 350000},
              'rep_asp_art': {'treat': 'maj', 'mat': 'asp', 'class': 'art',
                              'hci': 0, 'cost': 1000000},
              'min_asp_col': {'treat': 'min', 'mat': 'asp', 'class': 'col',
                              'hci': 85, 'cost': 100000},
              'maj_asp_col': {'treat': 'maj', 'mat': 'asp', 'class': 'col',
                              'hci': 60, 'cost': 350000},
              'rep_asp_col': {'treat': 'maj', 'mat': 'asp', 'class': 'col',
                              'hci': 0, 'cost': 1000000},
              'min_chip_col': {'treat': 'min', 'mat': 'chip', 'class': 'col',
                               'hci': 70, 'cost': 30000},
              'maj_chip_col': {'treat': 'maj', 'mat': 'chip', 'class': 'col',
                               'hci': 50, 'cost': 80000},
              'rep_chip_col': {'treat': 'maj', 'mat': 'chip', 'class': 'col',
                               'hci': 0, 'cost': 100000},
              'min_chip_loc': {'treat': 'min', 'mat': 'chip', 'class': 'loc',
                               'hci': 70, 'cost': 30000},
              'maj_chip_loc': {'treat': 'maj', 'mat': 'chip', 'class': 'loc',
                               'hci': 45, 'cost': 80000},
              'rep_chip_loc': {'treat': 'maj', 'mat': 'chip', 'class': 'loc',
                               'hci': 0, 'cost': 100000},
              }

ndf = pd.DataFrame(treatments)
ndf = ndf.round(3)

dft = ndf.T


def get_treatment(material, func_class, hci):
    return dft[(dft['mat'] == material) &
               (dft['class'] == func_class) &
               (dft['hci'].astype(int) >= hci)].index


test_data = {'a': {'mat': 'asp', 'class': 'art', 'hci': 35},
             'b': {'mat': 'asp', 'class': 'art', 'hci': 90},
             'c': {'mat': 'asp', 'class': 'col', 'hci': 68},
             'd': {'mat': 'asp', 'class': 'art', 'hci': 5},
             'e': {'mat': 'chip', 'class': 'col', 'hci': 55},
             'f': {'mat': 'chip', 'class': 'loc', 'hci': 95},
             'g': {'mat': 'asp', 'class': 'art', 'hci': 15},
             'h': {'mat': 'asp', 'class': 'art', 'hci': 70},
             'i': {'mat': 'asp', 'class': 'art', 'hci': 3}
             }

for k, v in test_data.items():
    res = get_treatment(v['mat'], v['class'], v['hci'])
    print (k, res)


# Output:
# a Index(['min_asp_art', 'maj_asp_art'], dtype='object')
# b Index([], dtype='object')
# c Index(['min_asp_col'], dtype='object')
# d Index(['min_asp_art', 'maj_asp_art'], dtype='object')
# e Index(['min_chip_col'], dtype='object')
# f Index([], dtype='object')
# g Index(['min_asp_art', 'maj_asp_art'], dtype='object')
# h Index(['min_asp_art'], dtype='object')
# i Index(['min_asp_art', 'maj_asp_art'], dtype='object')


from Pandas: interpolate between column values to return index name

No comments:

Post a Comment