How to apply function on very large pandas dataframe, with function depending on consecutive rows?

user_14168447 0
Viewed 583

I want to calculate the speed (m/s and km/h) with euclidian distance based on positions (x,y in meters) and time (in seconds). I found a way to take into account the fact that each time a name appears for the first time in dataframe, the speed is equal to NaN.

Problem: my dataframe is so large (> 1.5 millions rows) that, when I run the code, it is not done after more than 2 hours...

The code works with a shorter dataframe, the problem seems to be the length of the initial df.

Here is the simplified dataframe, followed by the code:

df
   name   time      x     y 
0  Mary      0     17    15
1  Mary      1   18.5    16
2  Mary      2     21    18
3  Steve     0     12    16
4  Steve     1   10.5    14
5  Steve     2      8    13
6  Jane      0     15    16
7  Jane      1     17    17
8  Jane      2     18    19

# calculating speeds:
for i in range(len(df)):
  if i >= 1:
    df.loc[i,'speed (m/s)'] = sqrt( (df.loc[i,'x'] - df.loc[i-1,'x'])**2 + (df.loc[i,'y'] - df.loc[i-1,'y'])**2 )
    df.loc[i,'speed (km/h)'] = df.loc[i,'speed (m/s)']*3.6

# each first time a name appears, speeds are equal to NaN:
first_indexes = []
names = df['name'].unique()

for j in names:
  a = df.index[df['name'] == j].tolist()
  if len(a) > 0 :
    first_indexes.append(a[0])

for index in first_indexes:
  df.loc[index, 'speed (m/s)'] = np.nan
  df.loc[index, 'speed (km/h)'] = np.nan

Iterating over this dataframe is way too long, I'm looking for a way to do this faster...

Thanks by advance for helping !

EDIT

df = pd.DataFrame([["Mary",0,17,15],
["Mary",1,18.5,16],
["Mary",2,21,18],
["Steve",0,12,16],
["Steve",1,10.5,14],
["Steve",2,8,13],
["Jane",0,15,16],
["Jane",1,17,17],
["Jane",2,18,19]],columns = [ "name","time","x","y" ])
pandas dataframe python
3 Answers

You can apply method for all data without loops and then set missing value for first name rows (data has to be sorted by name):

df['speed (m/s)'] = (np.sqrt(df['x'].sub(df['x'].shift()).pow(2) + 
                             df['y'].sub(df['y'].shift()).pow(2)) )
df['speed (km/h)'] = df['speed (m/s)']*3.6

cols = ['speed (m/s)','speed (km/h)']
df[cols] = df[cols].mask(~df['name'].duplicated())
print (df)
    name  time     x   y  speed (m/s)  speed (km/h)
0   Mary     0  17.0  15          NaN           NaN
1   Mary     1  18.5  16     1.802776      6.489992
2   Mary     2  21.0  18     3.201562     11.525624
3  Steve     0  12.0  16          NaN           NaN
4  Steve     1  10.5  14     2.500000      9.000000
5  Steve     2   8.0  13     2.692582      9.693297
6   Jane     0  15.0  16          NaN           NaN
7   Jane     1  17.0  17     2.236068      8.049845
8   Jane     2  18.0  19     2.236068      8.049845
user_29010020

Try this:

df = pd.read_csv('data.csv')

def calculate_speed(s):
    return sqrt((s['dx'])**2 + (s['dy'])**2)

df = df.join(df.groupby('name')[['x','y']].diff().rename({'x':'dx', 'y':'dy'}, axis=1))
df['speed (m/s)'] = df.apply(calculate_speed, axis=1)
df['speed (km/h)'] = df['speed (m/s)']*3.6
print(df)
user_108157180

If you want to work with a MultiIndex (which has many nice properties when working with dataframes with names and time indices), you could pivot your table to make name, x and y a column MultiIndex with time being the index:

dfp = df.pivot(index='time', columns=['name'])

Then you can easily calculate the speed for each name without having to check for np.NaN, duplicates or other invalid values:

speed_ms = np.sqrt((dfp['x'] - dfp['x'].shift(-1))**2 + (dfp['y'] - dfp['y'].shift(-1))**2).shift(1)

Now get the speed in km/h

speed_kmh = speed_ms * 3.6

And make both to a multiindex to make merging/concatenating the dataframes more explicit:

speed_ms.columns = pd.MultiIndex.from_product((['speed (m/s)'], speed_ms.columns))
speed_kmh.columns = pd.MultiIndex.from_product((['speed (km/h)'], speed_kmh.columns))

And finally concatenate the results to the dataframe. swaplevel makes all columns primarily indexable by the name, while sort_index sorts by the names:

dfp = pd.concat((dfp, speed_ms, speed_kmh), axis=1).swaplevel(1, 0, 1).sort_index(axis=1)

Now your dataframe looks like:

# Out[100]: 
name         Jane                        ...        Steve                      
     speed (km/h) speed (m/s)     x   y  ... speed (km/h) speed (m/s)     x   y
time                                     ...                                   
0             NaN         NaN  15.0  16  ...          NaN         NaN  12.0  16
1        8.049845    2.236068  17.0  17  ...     9.000000    2.500000  10.5  14
2        8.049845    2.236068  18.0  19  ...     9.693297    2.692582   8.0  13

[3 rows x 12 columns]

And you can easily index speeds and positions by names:

dfp['Mary']

#Out[107]: 
      speed (km/h)  speed (m/s)     x   y
time                                     
0              NaN          NaN  17.0  15
1         6.489992     1.802776  18.5  16
2        11.525624     3.201562  21.0  18

With dfp.stack(0) you re-transform it to your input-df-style, while keeping the names as a second index level:

dfp.stack(0).sort_index(level=1)

# Out[104]: 
            speed (km/h)  speed (m/s)     x   y
time name                                      
0    Jane            NaN          NaN  15.0  16
     Mary            NaN          NaN  17.0  15
     Steve           NaN          NaN  12.0  16
1    Jane       8.049845     2.236068  17.0  17
     Mary       6.489992     1.802776  18.5  16
     Steve      9.000000     2.500000  10.5  14
2    Jane       8.049845     2.236068  18.0  19
     Mary      11.525624     3.201562  21.0  18
     Steve      9.693297     2.692582   8.0  13

While dfp.stack(1) sets the names as columns, while setting speeds etc as indices.

user_63455180
Related