# Optimal way to filter DataFrame by many arguments, without usage of multiple ifs and for loops

I am learning about survival analysis and I made some implementation in Python using lifelines and plotly. However depending on the number of arguments I would have to almost rewrite my function again because of all these ifs and fors. Could you suggest the way I could make it smarter? Here is my code:

def classicLifelineFigure(frames, timeName = None, protocol = None, mutation = None, getter = getDfToSurvivalAnalysis):

'''

Function that returns lifeline figure object to be displayed on the website

Parameters:

frames --> dictionary with frames for features, mutations etc.
timeName --> time on which we perform analysis
protocol --> name of protocol if present in analysis
mutation --> name of mutation if present in analysis
getter --> function to get separate dataframes for analysis from one big dataframe

'''

kmf = KaplanMeierFitter()
fig = go.Figure()

if timeName == None:
return fig

dfToAnalysis = getter(frames, timeName, protocol, mutation)
dfToAnalysis['TIME'] = dfToAnalysis['TIME'] / 365

if protocol != None:

if mutation != None:

for prot in set(dfToAnalysis['PROTOCOL'].values):
for mut in set(dfToAnalysis['MUTATION'].values):

tempDf = dfToAnalysis[(dfToAnalysis['PROTOCOL'] == prot) & (dfToAnalysis['MUTATION'] == mut)]
results = kmf.fit(tempDf['TIME'], tempDf['EVENT'])
traceDf = results.survival_function_.copy()
x = traceDf.index,
y = traceDf.iloc[:,0],
mode = 'lines',
name = 'Protocol value = {}, mutation value = {}'.format(prot, mut)
))

if mutation == None:
for prot in set(dfToAnalysis['PROTOCOL'].values):
tempDf = dfToAnalysis[dfToAnalysis['PROTOCOL'] == prot]
results = kmf.fit(tempDf['TIME'], tempDf['EVENT'])
traceDf = results.survival_function_.copy()
x = traceDf.index,
y = traceDf.iloc[:,0],
mode = 'lines',
name = 'Protocol value = {}'.format(prot)
))

elif mutation != None:
for mut in set(dfToAnalysis['MUTATION'].values):

tempDf = dfToAnalysis[dfToAnalysis['MUTATION'] == mut]
results = kmf.fit(tempDf['TIME'], tempDf['EVENT'])
traceDf = results.survival_function_.copy()
x = traceDf.index,
y = traceDf.iloc[:,0],
mode = 'lines',
name = 'Mutation value = {}'.format(mut)
))

else:
results = kmf.fit(dfToAnalysis['TIME'], dfToAnalysis['EVENT'])
traceDf = results.survival_function_.copy()
x = traceDf.index,
y = traceDf.iloc[:,0],
mode = 'lines',
name = 'Survival function'
))

fig.update_layout(showlegend = True)

return fig

• Welcome to Code Review! Here on this site, complete code is preferred in order to see the program in context, i.e. that it's often better to post the code including all imports and maybe a small example on how the function is supposed to be used. – AlexV Aug 1 '19 at 8:49

Comparisons to None should always be done with is or is not. This is because None is a singleton object and this will give you the right behavior even if you override some objects __eq__ method or similar. This is mentioned in Python's official style-guide, PEP8. It also recommends no spaces around = for keyword arguments and using lower_case for variables and functions.

As to your actual question, what you really want to do is apply some function (fitting and plotting) to either the whole dataset or grouped by one or more categorical values. The latter is what pandas.DataFrame.groupby is for.

kmf = KaplanMeierFitter()

def plot_fit(df, fig, name='Survival function'):
results = kmf.fit(df['TIME'], df['EVENT'])
trace_df = results.survival_function_.copy()  # Is this copy really needed?
x=trace_df.index,
y=trace_df.iloc[:,0],
mode='lines',
name=name
))

def do_fit(df, by=None):
fig = go.Figure()
if by is None:
plot_fit(df, fig)
else:
for group, df2 in df.groupby(by):
name = ...  # Somehow determine label from group
plot_fit(df2, fig, name)
fig.update_layout(showlegend=True)
return fig


This function can be called either as do_fit(df), do_fit(df, "PROTOCOL") or do_fit(df, ["PROTOCOL", "MUTATION"]). The group variable is a tuple telling you which protocol and/or mutation it is, extracting a correct label from this is left as an exercise. I also omitted getting the dataframe with your getter function and other preparations. These should probably be done in the function calling this function.

It could be argued that this should be divided up even further and the fitting should be separated from the plotting.

• Wow, that's great way to do that. I figured out to use groupby for that one but this solution is great. Thank you! – Slajni Aug 6 '19 at 8:06