Apply function to sets of columns in pandas, 'looping' over entire data frame column-wise

Question

Here is a test example to show what I am trying to achieve. Here's a toy data frame:

df = pd.DataFrame(np.random.randn(10,7),index=range(1,11),columns=headers)

Which gives

    Time       A_x       A_y       A_z       B_x       B_y       B_z
1  -0.075509 -0.123527 -0.547239 -0.453707 -0.969796  0.248761  1.369613
2  -0.206369 -0.112098 -1.122609  0.218538 -0.878985  0.566872 -1.048862
3  -0.194552  0.818276 -1.563931  0.097377  1.641384 -0.766217 -1.482096
4   0.502731  0.766515 -0.650482 -0.087203 -0.089075  0.443969  0.354747
5   1.411380 -2.419204 -0.882383  0.005204 -0.204358 -0.999242 -0.395236
6   1.036695  1.115630  0.081825 -1.038442  0.515798 -0.060016  2.669702
7   0.392943  0.226386  0.039879  0.732611 -0.073447  1.164285  1.034357
8  -1.253264  0.389148  0.158289  0.440282 -1.195860  0.872064  0.906377
9  -0.133580 -0.308314 -0.839347 -0.517989  0.652120  0.477232 -0.391767
10  0.623841  0.473552  0.059428  0.726088 -0.593291 -3.186297 -0.846863

What I want to do is simply to calculate the length of the vector for each header (A and B) in this case, for each index, and divide by the Time column. Hence, this function needs to be np.sqrt(A_x^2 + A_y^2 + A_z^2) and the same for B of course. Ie I am looking to calculate the velocity for each row, but three columns contribute to one velocity result.

I have tried using df.groupby and df.filter to loop-over the columns but I cannot really get it to work, because I am not at all sure how I apply effectively the same function to chunks of the data-frame, all in one go (as apparently one is to avoid looping over rows). I have tried doing

df = df.apply(lambda x: np.sqrt(x.dot(x)), axis=1)

This works of course, but only if the input data frame has the right number of columns (3), if longer then the dot-product is calculated over the entire row, and not in chunks of three columns which is what I want (because this is turns corresponds to the tag coordinates, which are three dimensional).

So this is what I am eventually trying to get with the above example (the below arrays are just filled with random numbers, not the actual velocities which I am trying to calculate - just to show what sort of shape I trying to achieve):

     Velocity_A  Velocity_B
1    -0.975633   -2.669544
2     0.766405   -0.264904
3     0.425481   -0.429894
4    -0.437316    0.954006
5     1.073352   -1.475964
6    -0.647534    0.937035
7     0.082517    0.438112
8    -0.387111   -1.417930
9    -0.111011    1.068530
10    0.451979   -0.053333

My actual data is 50,000 x 36 (so there are 12 tags with x,y,z coordinates), and I want to calculate the velocity all in one go to avoid iterating (if at all possible). There is also a time column of the same length (50,000x1).

How do you do this?

Thanks, Astrid

Answer 1

A possible start.

Filtering out column names corresponding to a particular vector. For example

In [20]: filter(lambda x: x.startswith("A_"),df.columns)
Out[20]: ['A_x', 'A_y', 'A_z']

Sub selecting these columns from the DataFrame

In [22]: df[filter(lambda x: x.startswith("A_"),df.columns)]
Out[22]: 
         A_x       A_y       A_z
1  -0.123527 -0.547239 -0.453707
2  -0.112098 -1.122609  0.218538
3   0.818276 -1.563931  0.097377
4   0.766515 -0.650482 -0.087203
5  -2.419204 -0.882383  0.005204
6   1.115630  0.081825 -1.038442
7   0.226386  0.039879  0.732611
8   0.389148  0.158289  0.440282
9  -0.308314 -0.839347 -0.517989
10  0.473552  0.059428  0.726088

So, using this technique you can get chunks of 3 columns. For example.

column_initials = ["A","B"]
for column_initial in column_initials:
    df["Velocity_"+column_initial]=df[filter(lambda x: x.startswith(column_initial+"_"),df.columns)].apply(lambda x: np.sqrt(x.dot(x)), axis=1)/df.Time


In [32]: df[['Velocity_A','Velocity_B']]
Out[32]: 
    Velocity_A  Velocity_B
1    -9.555311  -22.467965
2    -5.568487   -7.177625
3    -9.086257  -12.030091
4     2.007230    1.144208
5     1.824531    0.775006
6     1.472305    2.623467
7     1.954044    3.967796
8    -0.485576   -1.384815
9    -7.736036   -6.722931
10    1.392823    5.369757

I do not get the same answer as yours. But, I borrowed your df.apply(lambda x: np.sqrt(x.dot(x)), axis=1) and assume it is correct.

Hope this helps.

Answer 2

Your calculation is more NumPy-ish than Panda-ish, by which I mean the calculation can be expressed somewhat succinctly if you regard your DataFrame as merely a big array, whereas the solution (at least the one I came up with) is more complicated when you try to wrangle the DataFrame with melt, groupby, etc.

The entire calculation can be expressed in essentially one line:

np.sqrt((arr**2).reshape(arr.shape[0],-1,3).sum(axis=-1))/times[:,None]

So here is the NumPy way:

import numpy as np
import pandas as pd
import io
content = '''
Time       A_x       A_y       A_z       B_x       B_y       B_z
-0.075509 -0.123527 -0.547239 -0.453707 -0.969796  0.248761  1.369613
-0.206369 -0.112098 -1.122609  0.218538 -0.878985  0.566872 -1.048862
-0.194552  0.818276 -1.563931  0.097377  1.641384 -0.766217 -1.482096
 0.502731  0.766515 -0.650482 -0.087203 -0.089075  0.443969  0.354747
 1.411380 -2.419204 -0.882383  0.005204 -0.204358 -0.999242 -0.395236
 1.036695  1.115630  0.081825 -1.038442  0.515798 -0.060016  2.669702
 0.392943  0.226386  0.039879  0.732611 -0.073447  1.164285  1.034357
-1.253264  0.389148  0.158289  0.440282 -1.195860  0.872064  0.906377
-0.133580 -0.308314 -0.839347 -0.517989  0.652120  0.477232 -0.391767
 0.623841  0.473552  0.059428  0.726088 -0.593291 -3.186297 -0.846863'''

df = pd.read_table(io.BytesIO(content), sep='\s+', header=True)

arr = df.values
times = arr[:,0]
arr = arr[:,1:]
result = np.sqrt((arr**2).reshape(arr.shape[0],-1,3).sum(axis=-1))/times[:,None]
result = pd.DataFrame(result, columns=['Velocity_%s'%(x,) for x in list('AB')])
print(result)

which yields

   Velocity_A  Velocity_B
0   -9.555311  -22.467965
1   -5.568487   -7.177625
2   -9.086257  -12.030091
3    2.007230    1.144208
4    1.824531    0.775006
5    1.472305    2.623467
6    1.954044    3.967796
7   -0.485576   -1.384815
8   -7.736036   -6.722931
9    1.392823    5.369757

---

Since your actual DataFrame has shape (50000, 36), choosing a quick method may be important. Here is a benchmark:

import numpy as np
import pandas as pd
import string

N = 12
col_ids = string.letters[:N]
df = pd.DataFrame(
    np.random.randn(50000, 3*N+1), 
    columns=['Time']+['{}_{}'.format(letter, coord) for letter in col_ids
                      for coord in list('xyz')])


def using_numpy(df):
    arr = df.values
    times = arr[:,0]
    arr = arr[:,1:]
    result = np.sqrt((arr**2).reshape(arr.shape[0],-1,3).sum(axis=-1))/times[:,None]
    result = pd.DataFrame(result, columns=['Velocity_%s'%(x,) for x in col_ids])
    return result

def using_loop(df):
    results = pd.DataFrame(index=df.index) # the result container
    for id in col_ids:
        results['Velocity_'+id] = np.sqrt((df.filter(regex=id+'_')**2).sum(axis=1))/df.Time
    return results

Using IPython :

In [43]: %timeit using_numpy(df)
10 loops, best of 3: 34.7 ms per loop

In [44]: %timeit using_loop(df)
10 loops, best of 3: 82 ms per loop

Answer 3

I would do at least a loop over the tag identifier, but don't worry, that's a very fast loop that just determines the filter pattern to get the right columns:

df = pd.DataFrame(np.random.randn(10,7), index=range(1,11), columns='Time A_x A_y A_z B_x B_y B_z'.split())

col_ids = ['A', 'B'] # I guess you can create that one easily

results = pd.DataFrame(index=df.index) # the result container

for id in col_ids:
    results['Velocity_'+id] = np.sqrt((df.filter(regex=id+'_')**2).sum(axis=1))/df.Time

Answer 4

One liner...split over many lines for readability:

import numpy as np
import pandas as pd

np.random.seed(0)

df = pd.DataFrame(
        np.random.randn(10,7), 
        index=range(1,11), 
        columns='Time A_x A_y A_z B_x B_y B_z'.split()
        )

result = df\
    .loc[:, df.columns.values!='Time']\
    .T\
    .groupby(lambda x: x[0])\
    .apply(lambda x: np.sqrt((x ** 2).sum()))\
    .T\
    .apply(lambda x: x / df['Time'])

print result

            A          B
1    1.404626   1.310639
2   -2.954644 -10.874091
3    3.479836   6.105961
4    3.885530   2.244544
5    0.995012   1.434228
6   11.278208  11.454466
7   -1.209242  -1.281165
8   -5.175911  -5.905070
9   11.889318  16.758958
10  -0.978014  -0.590767

Note: I am a bit frustrated that I needed to thrown in the two transposes. I just couldn't get groupby and apply to play nicely with axis=1 . If someone could show me how to do that, I'd be very grateful. The trick here was knowing that when you call groupby(lambda x: f(x)) that x is the value of the index for each row. So groupby(lambda x: x[0]) groups by the first letter of the row index. After doing the transposition, this was A or B .

---

Ok, no more transposes:

result = df\
    .loc[:, df.columns!='Time']\
    .groupby(lambda x: x[0], axis=1)\
    .apply(lambda x: np.sqrt((x**2).sum(1)))\
    .apply(lambda x: x / df['Time'])

print result

            A          B
1    1.404626   1.310639
2   -2.954644 -10.874091
3    3.479836   6.105961
4    3.885530   2.244544
5    0.995012   1.434228
6   11.278208  11.454466
7   -1.209242  -1.281165
8   -5.175911  -5.905070
9   11.889318  16.758958
10  -0.978014  -0.590767