Handling errors with pandas groupby and removing conflicting rows

Viewed 48

I have a dataset that contains population data, I have to compute the difference between two years. I achieve this with the following code:

data = [
    {'id': '1', 'year': 2010, 'total_pob': 100, 'total_pob_fem': 20},
    {'id': '1', 'year': 2020, 'total_pob': 200, 'total_pob_fem': 50},
    {'id': '2', 'year': 2010, 'total_pob': 150, 'total_pob_fem': 50},
    {'id': '2', 'year': 2020, 'total_pob': 300, 'total_pob_fem': 100},
    {'id': '3', 'year': 2010, 'total_pob': 340, 'total_pob_fem': 25},
    {'id': '3', 'year': 2020, 'total_pob': 400, 'total_pob_fem': 50},
    {'id': '4', 'year': 2010, 'total_pob': 50, 'total_pob_fem': 60},
    {'id': '4', 'year': 2020, 'total_pob': 100, 'total_pob_fem': 150},
]

df = pd.DataFrame(data)

pd.concat([df.set_index(['id', 'year'])
  .groupby('id')
  .apply(lambda x: x.iloc[0]-x.iloc[1])
  .assign(year='diff')
  .set_index('year', append=True),
  df.set_index(['id', 'year'])]).reset_index()

With the following output:

id year total_pob total_pob_fem
1 diff -100 -30
2 diff -150 -50
3 diff -60 -25
4 diff -50 -90
1 2010 100 20
1 2020 200 50
2 2010 150 50
2 2020 300 100
3 2010 340 25
3 2020 400 50
4 2010 50 60
4 2020 100 150

But when I try to compute the same difference, with data that was only one sample per year, it fails because it can't subtract it.

data = [
    {'id': '1', 'year': 2010, 'total_pob': 100, 'total_pob_fem': 20},
    {'id': '1', 'year': 2020, 'total_pob': 200, 'total_pob_fem': 50},
    {'id': '2', 'year': 2010, 'total_pob': 150, 'total_pob_fem': 50},
    {'id': '2', 'year': 2020, 'total_pob': 300, 'total_pob_fem': 100},
    {'id': '3', 'year': 2010, 'total_pob': 340, 'total_pob_fem': 25},
    {'id': '3', 'year': 2020, 'total_pob': 400, 'total_pob_fem': 50},
    # {'id': '4', 'year': 2010, 'total_pob': 50, 'total_pob_fem': 60},
    {'id': '4', 'year': 2020, 'total_pob': 100, 'total_pob_fem': 150},
]

or

data = [
    {'id': '1', 'year': 2010, 'total_pob': 100, 'total_pob_fem': 20},
    {'id': '1', 'year': 2020, 'total_pob': 200, 'total_pob_fem': 50},
    {'id': '2', 'year': 2010, 'total_pob': 150, 'total_pob_fem': 50},
    {'id': '2', 'year': 2020, 'total_pob': 300, 'total_pob_fem': 100},
    {'id': '3', 'year': 2010, 'total_pob': 340, 'total_pob_fem': 25},
    {'id': '3', 'year': 2020, 'total_pob': 400, 'total_pob_fem': 50},
    {'id': '4', 'year': 2010, 'total_pob': 50, 'total_pob_fem': 60},
    # {'id': '4', 'year': 2020, 'total_pob': 100, 'total_pob_fem': 150},
]

Is there a way to drop the rows that have only one data point per year, such as the output is:

id year total_pob total_pob_fem
1 diff -100 -30
2 diff -150 -50
3 diff -60 -25
1 2010 100 20
1 2020 200 50
2 2010 150 50
2 2020 300 100
3 2010 340 25
3 2020 400 50
4 2010 50 60
4 2020 100 150
1 Answers

I just can propose to calculate diff in another way and it will solve your problem. In your example you mixed int years with string 'diff' (and values in other columns) which complicates further processing. So, you can calculate diffs in new columns and you won't even meet such problems:

import pandas as pd

data = [
    {'id': '1', 'year': 2010, 'total_pob': 100, 'total_pob_fem': 20},
    {'id': '1', 'year': 2020, 'total_pob': 200, 'total_pob_fem': 50},
    {'id': '2', 'year': 2010, 'total_pob': 150, 'total_pob_fem': 50},
    {'id': '2', 'year': 2020, 'total_pob': 300, 'total_pob_fem': 100},
    {'id': '3', 'year': 2010, 'total_pob': 340, 'total_pob_fem': 25},
    {'id': '3', 'year': 2020, 'total_pob': 400, 'total_pob_fem': 50},
    {'id': '4', 'year': 2010, 'total_pob': 50, 'total_pob_fem': 60},
    # {'id': '4', 'year': 2020, 'total_pob': 100, 'total_pob_fem': 150},
]

df = pd.DataFrame(data)
df[['diff_total_pob','diff_total_pob_fem']] =  df.groupby('id').agg({'total_pob':'shift', 'total_pob_fem':'shift'}) - df[['total_pob','total_pob_fem']]

print(df)
>>
  id  year  total_pob  total_pob_fem  diff_total_pob  diff_total_pob_fem
0  1  2010        100             20             NaN                 NaN
1  1  2020        200             50          -100.0               -30.0
2  2  2010        150             50             NaN                 NaN
3  2  2020        300            100          -150.0               -50.0
4  3  2010        340             25             NaN                 NaN
5  3  2020        400             50           -60.0               -25.0
6  4  2010         50             60             NaN                 NaN
Related