Include an if statement when filtering a dataframe by other criteria

Question

I have a dataframe with about 600 rows, dating back to 2015:

ID   Year  Trap                             SVL
1    2020  ['P90', 'R81']                   78
2    2019  ['P3']                           56
3    2018  ['P22', 'R1', 'R11', 'PR22']     45

I need to search/filter individuals at hand (lizards), the known variables are the Trap and SVL (snout to vent length in mm) - for example, the current lizard might be from Trap == 'R1' and its SVL >= 45, then the search should return ID 3. In 2021, the lizard has an SVL = 75 (it grows by ~10mm/yr and I add +/- 5mm for measurement error).

I keep the criteria in a dict:

dict_choice = {'Trap': ['R1'],
               'SVL': (70.0, 80.0)}

My normal search is the following:

newdf = df.loc[df.Trap.map(set(dict_choice['Trap']).issubset) &
               ( (df['SVL'] >= dict_choice['SVL'][0]) & (df['SVL'] <= dict_choice['SVL'][1]))

Can I include the years as a search criteria? To minimise the number of matching lizards, I need to add a filter where for every year going back from 2021, we look through SVLs which are 5mm less. E.g lizard ID 3 was 45mm in 2018, which means that in 2021 it should be 75mm, however to match it we only need to look at lizars with SVL<=45. Something like this logic (if you can call it that!):

if @year == 2021: 
   search all where @Trap IN dict_choice['Trap'] AND ((df['SVL'] >= dict_choice['SVL'][0]) & (df['SVL'] <= dict_choice['SVL'][1]))
elif @year == 2020:
   search all where @Trap IN dict_choice['Trap'] AND ((df['SVL'] >= dict_choice['SVL'][0]) & (df['SVL'] <= dict_choice['SVL'][1]-5))
elif @year == 2019:
   search all where @Trap IN dict_choice['Trap'] AND ((df['SVL'] >= dict_choice['SVL'][0]) & (df['SVL'] <= dict_choice['SVL'][1]-10))
(...)

TL;DR: Can I include an if statement when filtering a dataframe by other criteria?

user_5478373 · Answer

Could I suggest that you add a "fixed" year to your test criteria like

dict_choice = {'Trap': ['R1'],
               'SVL': (70.0, 80.0),
               'Year': [2020]}

Then you can have the linear allowance for the test in the code itself

import pandas as pd
data = { 'ID': [1,2,3],
         'Year' : [2020,2019,2018],
         'Trap' : [['P90', 'R81'], ['P3'] ,['P22', 'R1', 'R11', 'PR22']],
         'SVL' : [78,56,45]}

df = pd.DataFrame(data)

gets the data frame as above

dict_choice = {'Trap': ['R81'],
               'SVL': (70.0, 80.0),
               'Year': [2018]}

is the enhanced query with a base year for the test

newdf = df.loc[df.Trap.map(set(dict_choice['Trap']).issubset) &
               ( (df['SVL'] >= dict_choice['SVL'][0]) & 
                 (df['SVL'] <= dict_choice['SVL'][1] + 5.0 * (df['Year'] - dict_choice['Year'])))]

Then this query has the linear allowance built into the check.

Note the query allows the lizard to grow if the year is ahead of the base year, and shink if it is behind.

Maybe you should have the movement on both sides but it just follows the pattern of the query (as below)

newdf = df.loc[df.Trap.map(set(dict_choice['Trap']).issubset) &
               ( (df['SVL'] >= dict_choice['SVL'][0] + 5.0 * (df['Year'] - dict_choice['Year'])) & 
                 (df['SVL'] <= dict_choice['SVL'][1] + 5.0 * (df['Year'] - dict_choice['Year'])))]

Include an if statement when filtering a dataframe by other criteria

1 Answers