python - pandas: complex filter on rows of DataFrame

Question

I would like to filter rows by a function of each row, e.g.

def f(row):
  return sin(row['velocity'])/np.prod(['masses']) > 5

df = pandas.DataFrame(...)
filtered = df[apply_to_all_rows(df, f)]

Or for another more complex, contrived example,

def g(row):
  if row['col1'].method1() == 1:
    val = row['col1'].method2() / row['col1'].method3(row['col3'], row['col4'])
  else:
    val = row['col2'].method5(row['col6'])
  return np.sin(val)

df = pandas.DataFrame(...)
filtered = df[apply_to_all_rows(df, g)]

How can I do so?

Answer 1

You can do this using DataFrame.apply, which applies a function along a given axis,

In [3]: df = pandas.DataFrame(np.random.randn(5, 3), columns=['a', 'b', 'c'])

In [4]: df
Out[4]: 
          a         b         c
0 -0.001968 -1.877945 -1.515674
1 -0.540628  0.793913 -0.983315
2 -1.313574  1.946410  0.826350
3  0.015763 -0.267860 -2.228350
4  0.563111  1.195459  0.343168

In [6]: df[df.apply(lambda x: x['b'] > x['c'], axis=1)]
Out[6]: 
          a         b         c
1 -0.540628  0.793913 -0.983315
2 -1.313574  1.946410  0.826350
3  0.015763 -0.267860 -2.228350
4  0.563111  1.195459  0.343168