How to choose stocks to invest in with Python | by Khuyen Tran | Towards Data Science

You plan to invest in several stocks in the coming 3 years, each with a different expected return for each dollar invested and a specific amount of investment, as shown in the table below (all in thousands of dollars). Given that the amount you can invest in stock purchases is limited each year, you want to decide which stocks to invest in each year so as to maximize the total returns.
Everybody wants to receive the most returns with the least investment but not everybody takes the time to analyze the stocks because this is not an easy task. You decided to use your data science skills to help you with this problem but found it challenging to solve with machine learning techniques. Why is this the case?
While Machine Learning is great in learning from data and making predictions, it has a limitation in making the optimal decision. But with Python MIP (Mixed-Integer Linear Programming) tool, the stocks chosen for each year could be determined with ease. When your problems involve multiple constraints and allow for many possible solutions, only MIP has the power to find the optimal solution. You can find more information about MIP here. This blog will show you how to use this tool for your benefit.
Input parameters:
Number of stocks (I)
Budget requirement of each stock for each year (a[t][i])
Budget requirement each year (b[t])
Annual return of each stock (c[i])
Decision variables: Whether or not to select a stock each year x[t][i]. This is a binary variable with x[t][i] = 1 when selected and x[t][i]=0 otherwise
Constraints: The sum of the budget requirement for all invested stocks each year cannot exceed the budget requirement for that year
Objective: Maximize total return in investment
c = [90, 120, 100, 80, 130]b = [45, 60, 50]a = [[10, 15, 12, 9, 13],[20, 15, 25, 15, 10],[15, 20, 20, 15, 10]]T, I = range(len(b)),range(len(c))
Install MIP
pip install mip==1.4.2 from mip.model import *
m = Model(sense = MAXIMIZE,solver_name=CBC)
Create variables with add_var()
Create a vector of variables with n binary decision variables (n = 5) indicating whether a stock is selected or not
x = [[m.add_var(name = names[i], var_type=BINARY) for i in I] for t in T]
Use *xsum()* for summation expressions
m.objective = maximize(xsum(c[i] * x[t][i] for i in I for t in T))
Add constraint in money available each year
for t in range(3):   m += xsum(a[t][i] * x[t][i] for i in I) <= b[t]
Optimize() method executes the optimization of a formulation. Optimize method returns the status (OptimizationStatus) of the search. OPTIMAL if the search was concluded and the optimal solution was found.
status = m.optimize() if status == OptimizationStatus.OPTIMAL:   print('optimal solution cost {} found'.format(m.objective_value)) 
results = {}for i in range(3):   results[i+1] = [m.vars[j].x for j in range(5*i,5*i+5)]
import pandas as pd result = pd.DataFrame(data = results, index = ['A','B','C','D','E'])
Based on the results obtained, this is what you should do to maximize the total return in investment with different limitation in the budget each year:
In year 1, invest in stock A, C, D, and E
In year 2, invest in stock A, B, and D
In year 3, invest in stock B, C, and E
Even though stocks E and B give the highest return on investment, they are not always chosen because of their high budget requirements in that respective year. (i.e Stock B is not chosen in year 1 because it has the highest budget requirement compared to other stocks)
Some stocks are not chosen in one year not because of their budget requirement but because of their low return in investment (i.e Stock A is not chosen in year 3 even though the budget requirement is relatively low)
Congratulations! You learned how which stocks to choose for each given year to achieve the optimal return in value using MIP. I hope you have fun choosing your stock and observing the optimal results after just a few lines of code.
Feel free to fork and play with the code for this article in this Github repo.
I like to write about basic data science concepts and play with different algorithms and data science tools. You could connect with me on LinkedIn and Twitter.
Star this repo if you want to check out the codes for all of the articles I have written. Follow me on Medium to stay informed with my latest data science articles like these:
towardsdatascience.com
towardsdatascience.com
towardsdatascience.com
medium.com
[1]Gurobi. Retrieved from: https://www.gurobi.com/machine-learning-and-mathematical-optimization/