Code that performs 2500 simulations by running a regression after randomly assigning 10 friends to each observation in a data took 20 days to run

Question

I am running the following code to conduct a simulation. I have a dataset of around 100K observations. The simulation does the following:

1. It randomly assign 10 "friends" to each observation
2. Take the average of the variable normalized_score of the 10 friends
3. Run a regression where the left hand side variable is the normalized_score and the left hand side variable is the average normalized_score of the 10 friends
4. Save the coefficient from the regression

I run this simulation 2500 times. It took around 20 days to finish. I tried parallelization to perform several simulations at once and randomization with replacement to make the code more efficient with no luck (the code crashed). I would like to ideally give each observation a 100 random friends, but I would like to make the code faster before doing this. Could you please provide some suggestions on how to make the iteration in the select_friends function faster and more efficient?

Here is the python code I wrote. It works well and produces the desired results. It just takes way too long.

I cannot share the data since it is not publicly available.

import pandas as pd
import numpy as np
from pyprojroot.here import here
import statsmodels.api as sm # for linear regression 
import random
random.seed(42)  # You can use any number as the seed

# Read the CSV file into a DataFrame
dir = "path"
data1 = pd.read_csv(dir + 'data1.csv')
data2 = pd.read_csv("path")

data = pd.merge(data1, data2, how='right', on='ID', suffixes=('', '_inSchool'))

def calculate_average_score(data, friends_dict):
    average_scores = []
    for key in friends_dict:
        friends = friends_dict[key]['friends']
        
        friends_scores = data[data['ID'].isin(friends)]['normalized_score'].mean()
        
        average_scores.append({
            'ID': key,
            'friends_avg_score': friends_scores
        })
    return pd.DataFrame(average_scores)

def select_friends(data):
    friends_dict = {}
    for index, row in data.iterrows():
        friends = data.sample(n=10, replace=False)
        friends_dict[row['ID']] = {
            'friends': friends['ID'].tolist()
        }
    return friends_dict

    coefficients = []


for i in range(2500):
    friends_dict = select_friends(data)
    avg_score_df = calculate_average_score(data, friends_dict)

    data = data.merge(avg_score_df, on='ID', how='left')

    formula = "normalized_score ~ 1 + friends_avg_score"
    ols_model = sm.OLS.from_formula(formula, data=data).fit()
    
    # Extract and store the coefficient of interest (friends_avg_score)
    coefficients.append(ols_model.params['friends_avg_score'])

Edit: Here is a code that would generate a sample that is somewhat representative

import pandas as pd
import numpy as np

# Set the random seed for reproducibility
np.random.seed(0)

# Create an ID column
id_column = range(1, 100001)

# Create a normalized_score column with random values from a normal distribution
# The standard normal distribution has a mean of 0 and a standard deviation of 1
# We'll use np.clip to ensure that the values lie within the range [-2, 2]
normalized_score_column = np.clip(np.random.randn(100000), -2, 2)

# Create a DataFrame
data = pd.DataFrame({'ID': id_column, 'normalized_score': normalized_score_column})

# If you want to save this DataFrame to a CSV file, you can use the following line:
# data.to_csv('sample_data.csv', index=False)

# Print the first few rows to check the data
print(data.head())
```

Answer 1

1. Avoid unnecessary operations:

   • In the select_friends() function, instead of sampling 10 friends for each student using sample(), you can directly select 10 random indices from the DataFrame data using random.sample() and retrieve the corresponding IDs. This will avoid repeatedly sampling from the DataFrame and improve performance.

2. Use vectorized operations:

   • Instead of using a loop in the calculate_average_score() function to calculate the average scores for each student, you can leverage vectorized operations in pandas to perform the calculation more efficiently. You can use the groupby() function to group the data by the student ID and then use the mean() function to calculate the average score for each group. This can significantly reduce the computation time.

3. Avoid unnecessary DataFrame merges:

   • Instead of merging the avg_score_df DataFrame with the data DataFrame in each iteration of the loop, you can store the average scores in a separate dictionary or DataFrame and update it incrementally. This can avoid the repeated merging operation, which can be computationally expensive.

4. Use parallel processing:

   • If your machine has multiple cores or you have access to a distributed computing environment, you can consider using parallel processing techniques to parallelize the computations. This can help to speed up the calculations by utilizing multiple processors or machines simultaneously.

Here's an optimized version of the code with the above suggestions implemented:

import pandas as pd
import numpy as np
from pyprojroot.here import here
import statsmodels.api as sm  # for linear regression
import random
from multiprocessing import Pool

random.seed(42)  # You can use any number as the seed

# Read the CSV file into a DataFrame
dir = "path"
data1 = pd.read_csv(dir + 'data1.csv')
data2 = pd.read_csv("path")

data = pd.merge(data1, data2, how='right', on='ID', suffixes=('', '_inSchool'))

def calculate_average_score(data, friends_dict):
    friends_scores = data[data['ID'].isin(friends_dict.keys())].groupby('ID')['normalized_score'].mean()
    return pd.DataFrame({'ID': list(friends_dict.keys()), 'friends_avg_score': friends_scores})

def select_friends(data):
    friends_dict = {}
    random_indices = random.sample(range(data.shape[0]), 10)
    selected_friends = data.loc[random_indices, 'ID'].tolist()
    for index, row in data.iterrows():
        friends_dict[row['ID']] = {
            'friends': selected_friends
        }
    return friends_dict

coefficients = []

def calculate_coefficient(seed):
    random.seed(seed)
    friends_dict = select_friends(data)
    avg_score_df = calculate_average_score(data, friends_dict)

    formula = "normalized_score ~ 1 + friends_avg_score"
    ols_model = sm.OLS.from_formula(formula, data=data.merge(avg_score_df, on='ID', how='left')).fit()

    return ols_model.params['friends_avg_score']

# Use parallel processing to calculate coefficients
with Pool() as pool:
    coefficients = pool.map(calculate_coefficient, range(2500))
```