Conditionally Create or Assign Columns on Pandas DataFrames | by Louis Chan | Towards Data Science

Python, arguably the coolest programming language these days (thanks for Machine Learning and Data Science), is not very well known for its efficiency when compared with one of the best programming language — C. An example of this would be conditional logics. When developing machine learning models, it is quite common that we would need to programmatically update labels based on hard-coded rules derived from statistical analysis or findings of the previous iteration. No shame in admitting it: I have been coding it with Pandas apply until one day I grew so fed up with the nested block that I decided to research (a.k.a Google) into alternative, more maintainable and more efficient ways of doing it (a.k.a basically any other ways of doing it)
The dataset that we are going to use is the iris dataset, which you would be able to get it for free via either pandas or seaborn.
Let’s say that after our initial analysis, we want to label the dataset with the following logic:
if sepal length < 5.1, then label 0;
otherwise if sepal width > 3.3 and sepal length < 5.8, then label 1;
otherwise if sepal width > 3.3 and petal length > 5.1, then label 2;
otherwise if sepal width > 3.3, petal length < 1.6 and either sepal length < 6.4 or petal width < 1.3, then label 3;
otherwise if sepal width > 3.3 and either sepal length < 6.4 or petal width < 1.3, then label 4;
otherwise if sepal width > 3.3, then label 5;
otherwise label 6
Before we dive into the code, let’s also quickly set a new label column to None:
iris['label'] = None
If you are still using this, this blog post is definitely the right place for you!
1min 29s ± 8.91 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
And the time it takes to run... Okay, let’s move on...
Pandas .apply(), straightforward, is used to apply a function along an axis of the DataFrame or on values of Series. For example, if we have a function f that sum an iterable of numbers (i.e. can be a list, np.array, tuple, etc.), and pass it to a dataframe like below, we will be summing across a row:
def f(numbers):    return sum(numbers)df['Row Subtotal'] = df.apply(f, axis=1)
In the above snippet, axis=1 indicates the direction of applying the function. .apply() would by default has axis=0, i.e. apply the function column by column; while axis=1 would apply the function row by row.
Now that we have a basic understanding of pandas .apply(), let’s code the logic for assigning the labels and see how long does it run for:
1.43 s ± 115 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1.43s for 150,000 rows is a very big improvement from the previous level, but is still definitely very slow. Imagine if you need to handle a dataset of millions of say transaction data or credit approvals, each time we want to apply a set of rules and assign the engineered feature to a column, it would take over 14s. Run it for enough columns, and you can slack and chill for the rest of the afternoon.
If you are familiar with SQL, assigning values to a new column with .loc[] is literally just an UPDATE statement with a dragon of WHERE conditions. So surely this should scale a lot better than applying a function to each of the rows or columns.
13.3 ms ± 837 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Now we are spending only one-tenth of the time compared to the previous, which means you have less excuse to leave to desk and start netflix when you are working from home. However, we have only been using pandas built in functions for now. Although pandas gives us a very convenient high level interface to interact with data tables, efficiency may be lost through layers of abstractions.
Numpy has a lower level interface which allows more efficient interaction with n-dimensional iterables (i.e. vectors, matrix, tensors...). Its methods are more than often C-based and uses optimised algorithms when it comes to more complex computations, making it way faster than our reinvented wheels. According to numpy’s official documentation, np.where() accepts the following syntax:
np.where(condition, return value if True, return value if False)
In essence, this is a dichotomous logic where a conditional will be evaluated as a boolean and return a value accordingly. The trick under the sleeve here is that the condition can actually be an iterable (i.e. a boolean ndarray type). This means that it is perfectly fine for us to pass in df['feature'] == 1 as the condition, and code the where logic as:
np.where(    df['feature'] == 1,     'It is one',     'It is not one')
So you may ask, how can we implement the logic we state above with a bisection function like np.where()? The answer is simple, yet disturbing. Nesting np.where()... (Spoiler alert: this snippet might be triggering)
3.6 ms ± 149 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Congrats, you have survived it. I cannot tell you how many times I have spent on counting the closing parentheses, but hey, this gets the job done! We have slashed another 10ms from pandas .loc[]. However, this snippet is simply not maintainable, which means, it is not acceptable.
Numpy .select(), the function that is intended to implement a multichotomous logic, unlike .where().
np.select(condlist, choicelist, default=0)
It uses a simular syntax as np.where(), except that the first argument is now a list of conditions, which should have the same length as the choices. One thing to bear in mind when using np.select() is that a choice will be selected as soon as the first condition has been met. This means that if a superset rule comes before a subset rule in the list, the subset choice will never be chosen. To put it in context:
condlist = [    df['A'] <= 1,    df['A'] < 1]choicelist = ['<=1', '<1']selection = np.select(condlist, choicelist, default='>1')
Because all of the rows that hits df['A'] < 1 will also be captured bydf['A'] <= 1, no rows will end up being labelled as '<1'. To avoid this from happening, please do make sure to have a less specific rule before the more specific ones:
condlist = [    df['A'] < 1, # < ───┬ swapped here    df['A'] <= 1 # < ───┘]choicelist = ['<1', '<=1'] # remember to update this as well!!selection = np.select(condlist, choicelist, default='>1')
As you can see from the above, you will need to update both condlist and choicelsit to make sure that the code runs smoothly. But seriously, who has time for that when Netflix is within our hand’s reach? Well, I have got you, pal! By changing it to a dictionary, we will be hitting roughly the same time and memory complexity but with a way more maintainable code snippet:
6.29 ms ± 475 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
About double that of nested np.where(), but not only does this save you from the table flipping parenthesis debugging, but also the absentminded choicelist change. I will be the first to bite the bullet, I have forgotten to update choicelist for way too many time that I have ended up spending more than quadruple the trying to debug my machine learning models. Trust me, np.select() with dict. It is worth it!
Numpy’s vectorised operations: If your code involves looping and evaluating unary functions, binary functions or functions that operate on sequences of numbers. You should definitely refactor your code by transforming your data into a numpy ndarray, and make full use of numpy’s vectorised operations to immensely speed up your script. Conclusion: This is a situational winner! Check out here at Numpy’s official documentation for examples of unary functions, binary functions or functions that operate on sequences of numbers!np.vectorize: Do not be fooled by the name of this function. This is just a convenience function and doesn't actually make code run any faster. To use this function, you will need to first code your logic into a callable function, then run np.vectorize(your_function)(your_data_series). Another big down side is that you would need to transform your dataframe into an one-dimensional iterable in order to pass into the “vectorised” function. Conclusion: if it’s not convenient to use np.vectorize, don’t even think about it.numba.njit: Now this is real business, real vectorisation. When it comes to optimised loops, you simply cannot easily beat numba without super genius hacky moves. It tries to move any numpy evaluation to as close to C as possible to supercharge its efficiency. While it can light-speed your numeric calculation, it also restricts itself to numeric calculations, which means no pandas series, no string indexing, just numpy’s ndarray with types int, float, datetime, bool, and category. Conclusion: A real CHAMPION if you are comfortable with using only Numpy’s ndarray and translating your logic into numeric computations and numeric computations only. Learn more from here.
Numpy’s vectorised operations: If your code involves looping and evaluating unary functions, binary functions or functions that operate on sequences of numbers. You should definitely refactor your code by transforming your data into a numpy ndarray, and make full use of numpy’s vectorised operations to immensely speed up your script. Conclusion: This is a situational winner! Check out here at Numpy’s official documentation for examples of unary functions, binary functions or functions that operate on sequences of numbers!
np.vectorize: Do not be fooled by the name of this function. This is just a convenience function and doesn't actually make code run any faster. To use this function, you will need to first code your logic into a callable function, then run np.vectorize(your_function)(your_data_series). Another big down side is that you would need to transform your dataframe into an one-dimensional iterable in order to pass into the “vectorised” function. Conclusion: if it’s not convenient to use np.vectorize, don’t even think about it.
numba.njit: Now this is real business, real vectorisation. When it comes to optimised loops, you simply cannot easily beat numba without super genius hacky moves. It tries to move any numpy evaluation to as close to C as possible to supercharge its efficiency. While it can light-speed your numeric calculation, it also restricts itself to numeric calculations, which means no pandas series, no string indexing, just numpy’s ndarray with types int, float, datetime, bool, and category. Conclusion: A real CHAMPION if you are comfortable with using only Numpy’s ndarray and translating your logic into numeric computations and numeric computations only. Learn more from here.
If you are interested in other Python tricks, I have put together a list of these short blogs for you:
Python Tricks: Flattening Lists
Python Tricks: How to Check Table Merging with Pandas
Python Tricks: Simplifying If Statements & Boolean Evaluation
Python Tricks: Check Multiple Variables against Single Value
If you want to read more about Python, Data Science, or Machine Learning, you may want to check out these posts:
7 Easy Ways for Improving Your Data Science Workflow
Efficient Conditional Logic on Pandas DataFrames
Memory Efficiency of Common Python Data Structures
Parallelism with Python
Essential Jupyter Extension for Data Science Set Up
Efficient Root Searching Algorithms in Python
If you would like to read more about how to apply machine learning to trading & investing, here are some other posts that may be of interest:
Genetic Algorithm for Trading Strategy Optimization in Python
Genetic Algorithm — Stop Overfitting Trading Strategies
ANN Recommendation System for Stock Selection
If possible, go for numba.njit; otherwise, np.select with a dict will help you sail afar, my friends. Remember, every little bit of improvement helps! Let me know if you have learnt something new from this! And please also let me know if there are other neat tricks that I have missed!